hotshot_example_types/membership/
strict_membership.rs

1use std::{collections::HashSet, fmt, fmt::Debug, sync::Arc};
2
3use alloy::primitives::U256;
4use async_broadcast::Receiver;
5use async_lock::RwLock;
6use hotshot_types::{
7    data::Leaf2,
8    drb::DrbResult,
9    event::Event,
10    stake_table::HSStakeTable,
11    traits::{
12        block_contents::BlockHeader,
13        election::{Membership, NoStakeTableHash},
14        node_implementation::{ConsensusTime, NodeImplementation, NodeType},
15        signature_key::StakeTableEntryType,
16    },
17    utils::{epoch_from_block_number, root_block_in_epoch, transition_block_for_epoch},
18    PeerConfig,
19};
20
21use crate::{
22    membership::{fetcher::Leaf2Fetcher, stake_table::TestStakeTable},
23    storage_types::TestStorage,
24};
25
26#[derive(Clone)]
27pub struct StrictMembership<
28    TYPES: NodeType,
29    StakeTable: TestStakeTable<TYPES::SignatureKey, TYPES::StateSignatureKey>,
30> {
31    inner: StakeTable,
32    epochs: HashSet<TYPES::Epoch>,
33    drbs: HashSet<TYPES::Epoch>,
34    fetcher: Arc<RwLock<Leaf2Fetcher<TYPES>>>,
35    epoch_height: u64,
36}
37
38impl<TYPES, StakeTable> Debug for StrictMembership<TYPES, StakeTable>
39where
40    TYPES: NodeType,
41    StakeTable: TestStakeTable<TYPES::SignatureKey, TYPES::StateSignatureKey>,
42{
43    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
44        f.debug_struct("StrictMembership")
45            .field("inner", &self.inner)
46            .field("epochs", &self.epochs)
47            .field("drbs", &self.drbs)
48            .finish()
49    }
50}
51
52impl<
53        TYPES: NodeType,
54        StakeTable: TestStakeTable<TYPES::SignatureKey, TYPES::StateSignatureKey>,
55    > StrictMembership<TYPES, StakeTable>
56{
57    fn assert_has_stake_table(&self, epoch: Option<TYPES::Epoch>) {
58        let Some(epoch) = epoch else {
59            return;
60        };
61        assert!(
62            self.epochs.contains(&epoch),
63            "Failed stake table check for epoch {epoch}"
64        );
65    }
66    fn assert_has_randomized_stake_table(&self, epoch: Option<TYPES::Epoch>) {
67        let Some(epoch) = epoch else {
68            return;
69        };
70        assert!(
71            self.drbs.contains(&epoch),
72            "Failed drb check for epoch {epoch}"
73        );
74    }
75}
76
77impl<
78        TYPES: NodeType,
79        StakeTable: TestStakeTable<TYPES::SignatureKey, TYPES::StateSignatureKey>,
80    > Membership<TYPES> for StrictMembership<TYPES, StakeTable>
81{
82    type Error = anyhow::Error;
83    type StakeTableHash = NoStakeTableHash;
84    type Storage = TestStorage<TYPES>;
85
86    fn new<I: NodeImplementation<TYPES>>(
87        quorum_members: Vec<hotshot_types::PeerConfig<TYPES>>,
88        da_members: Vec<hotshot_types::PeerConfig<TYPES>>,
89        storage: Self::Storage,
90        network: Arc<<I as NodeImplementation<TYPES>>::Network>,
91        public_key: TYPES::SignatureKey,
92        epoch_height: u64,
93    ) -> Self {
94        let fetcher = Leaf2Fetcher::new::<I>(network, storage, public_key);
95
96        Self {
97            inner: TestStakeTable::new(
98                quorum_members.into_iter().map(Into::into).collect(),
99                da_members.into_iter().map(Into::into).collect(),
100            ),
101            epochs: HashSet::new(),
102            drbs: HashSet::new(),
103            fetcher: RwLock::new(fetcher).into(),
104            epoch_height,
105        }
106    }
107
108    async fn set_external_channel(&mut self, external_channel: Receiver<Event<TYPES>>) {
109        self.fetcher
110            .write()
111            .await
112            .set_external_channel(external_channel)
113    }
114
115    fn stake_table(&self, epoch: Option<TYPES::Epoch>) -> HSStakeTable<TYPES> {
116        self.assert_has_stake_table(epoch);
117        let peer_configs = self
118            .inner
119            .stake_table(epoch.map(|e| *e))
120            .into_iter()
121            .map(Into::into)
122            .collect();
123        HSStakeTable(peer_configs)
124    }
125
126    fn da_stake_table(&self, epoch: Option<TYPES::Epoch>) -> HSStakeTable<TYPES> {
127        self.assert_has_stake_table(epoch);
128        let peer_configs = self
129            .inner
130            .da_stake_table(epoch.map(|e| *e))
131            .into_iter()
132            .map(Into::into)
133            .collect();
134        HSStakeTable(peer_configs)
135    }
136
137    fn committee_members(
138        &self,
139        _view_number: TYPES::View,
140        epoch: Option<TYPES::Epoch>,
141    ) -> std::collections::BTreeSet<TYPES::SignatureKey> {
142        self.assert_has_stake_table(epoch);
143        self.inner
144            .stake_table(epoch.map(|e| *e))
145            .into_iter()
146            .map(|entry| entry.signature_key)
147            .collect()
148    }
149
150    fn da_committee_members(
151        &self,
152        _view_number: TYPES::View,
153        epoch: Option<TYPES::Epoch>,
154    ) -> std::collections::BTreeSet<TYPES::SignatureKey> {
155        self.assert_has_stake_table(epoch);
156        self.inner
157            .da_stake_table(epoch.map(|e| *e))
158            .into_iter()
159            .map(|entry| entry.signature_key)
160            .collect()
161    }
162
163    fn stake(
164        &self,
165        pub_key: &TYPES::SignatureKey,
166        epoch: Option<TYPES::Epoch>,
167    ) -> Option<hotshot_types::PeerConfig<TYPES>> {
168        self.assert_has_stake_table(epoch);
169        self.inner
170            .stake(pub_key.clone(), epoch.map(|e| *e))
171            .map(Into::into)
172    }
173
174    fn da_stake(
175        &self,
176        pub_key: &TYPES::SignatureKey,
177        epoch: Option<TYPES::Epoch>,
178    ) -> Option<hotshot_types::PeerConfig<TYPES>> {
179        self.assert_has_stake_table(epoch);
180        self.inner
181            .da_stake(pub_key.clone(), epoch.map(|e| *e))
182            .map(Into::into)
183    }
184
185    /// Check if a node has stake in the committee
186    fn has_stake(
187        &self,
188        pub_key: &<TYPES as NodeType>::SignatureKey,
189        epoch: Option<<TYPES as NodeType>::Epoch>,
190    ) -> bool {
191        self.assert_has_stake_table(epoch);
192
193        self.stake(pub_key, epoch)
194            .is_some_and(|x| x.stake_table_entry.stake() > U256::ZERO)
195    }
196
197    /// Check if a node has stake in the da committee
198    fn has_da_stake(
199        &self,
200        pub_key: &<TYPES as NodeType>::SignatureKey,
201        epoch: Option<<TYPES as NodeType>::Epoch>,
202    ) -> bool {
203        self.assert_has_stake_table(epoch);
204
205        self.da_stake(pub_key, epoch)
206            .is_some_and(|x| x.stake_table_entry.stake() > U256::ZERO)
207    }
208
209    fn lookup_leader(
210        &self,
211        view: TYPES::View,
212        epoch: Option<TYPES::Epoch>,
213    ) -> anyhow::Result<TYPES::SignatureKey> {
214        self.assert_has_randomized_stake_table(epoch);
215        self.inner.lookup_leader(*view, epoch.map(|e| *e))
216    }
217
218    fn total_nodes(&self, epoch: Option<TYPES::Epoch>) -> usize {
219        self.assert_has_stake_table(epoch);
220        self.inner.stake_table(epoch.map(|e| *e)).len()
221    }
222
223    fn da_total_nodes(&self, epoch: Option<TYPES::Epoch>) -> usize {
224        self.assert_has_stake_table(epoch);
225        self.inner.da_stake_table(epoch.map(|e| *e)).len()
226    }
227
228    fn has_stake_table(&self, epoch: TYPES::Epoch) -> bool {
229        let has_stake_table = self.inner.has_stake_table(*epoch);
230
231        assert_eq!(has_stake_table, self.epochs.contains(&epoch));
232
233        has_stake_table
234    }
235
236    fn has_randomized_stake_table(&self, epoch: TYPES::Epoch) -> anyhow::Result<bool> {
237        if !self.has_stake_table(epoch) {
238            return Ok(false);
239        }
240        let has_randomized_stake_table = self.inner.has_randomized_stake_table(*epoch);
241
242        if let Ok(result) = has_randomized_stake_table {
243            assert_eq!(result, self.drbs.contains(&epoch));
244        } else {
245            assert!(!self.drbs.contains(&epoch));
246        }
247
248        has_randomized_stake_table
249    }
250
251    fn add_drb_result(&mut self, epoch: TYPES::Epoch, drb_result: hotshot_types::drb::DrbResult) {
252        self.assert_has_stake_table(Some(epoch));
253
254        self.drbs.insert(epoch);
255        self.inner.add_drb_result(*epoch, drb_result);
256    }
257
258    fn first_epoch(&self) -> Option<TYPES::Epoch> {
259        self.inner.first_epoch().map(TYPES::Epoch::new)
260    }
261
262    fn set_first_epoch(&mut self, epoch: TYPES::Epoch, initial_drb_result: DrbResult) {
263        self.epochs.insert(epoch);
264        self.epochs.insert(epoch + 1);
265
266        self.drbs.insert(epoch);
267        self.drbs.insert(epoch + 1);
268
269        self.inner.set_first_epoch(*epoch, initial_drb_result);
270    }
271
272    async fn add_epoch_root(
273        membership: Arc<RwLock<Self>>,
274        block_header: TYPES::BlockHeader,
275    ) -> anyhow::Result<()> {
276        let mut membership_writer = membership.write().await;
277
278        let epoch =
279            epoch_from_block_number(block_header.block_number(), membership_writer.epoch_height)
280                + 2;
281
282        membership_writer.epochs.insert(TYPES::Epoch::new(epoch));
283
284        membership_writer.inner.add_epoch_root(epoch);
285
286        Ok(())
287    }
288
289    async fn get_epoch_root(
290        membership: Arc<RwLock<Self>>,
291        epoch: TYPES::Epoch,
292    ) -> anyhow::Result<Leaf2<TYPES>> {
293        let membership_reader = membership.read().await;
294
295        let block_height = root_block_in_epoch(*epoch, membership_reader.epoch_height);
296
297        let stake_table = membership_reader.inner.stake_table(Some(*epoch));
298        let fetcher = membership_reader.fetcher.clone();
299
300        drop(membership_reader);
301
302        for node in stake_table {
303            if let Ok(leaf) = fetcher
304                .read()
305                .await
306                .fetch_leaf(block_height, node.signature_key)
307                .await
308            {
309                return Ok(leaf);
310            }
311        }
312
313        anyhow::bail!("Failed to fetch epoch root from any peer");
314    }
315
316    async fn get_epoch_drb(
317        membership: Arc<RwLock<Self>>,
318        epoch: TYPES::Epoch,
319    ) -> anyhow::Result<DrbResult> {
320        let membership_reader = membership.read().await;
321
322        let epoch_height = membership_reader.epoch_height;
323        let epoch_drb = membership_reader.inner.get_epoch_drb(*epoch);
324        let fetcher = membership_reader.fetcher.clone();
325
326        drop(membership_reader);
327
328        if let Ok(drb_result) = epoch_drb {
329            Ok(drb_result)
330        } else {
331            let previous_epoch = match epoch.checked_sub(1) {
332                Some(epoch) => epoch,
333                None => {
334                    anyhow::bail!("Missing initial DRB result for epoch {epoch:?}");
335                },
336            };
337
338            let drb_block_height = transition_block_for_epoch(previous_epoch, epoch_height);
339
340            let membership_reader = membership.read().await;
341            let stake_table = membership_reader.inner.stake_table(Some(previous_epoch));
342            drop(membership_reader);
343
344            let mut drb_leaf = None;
345
346            for node in stake_table {
347                if let Ok(leaf) = fetcher
348                    .read()
349                    .await
350                    .fetch_leaf(drb_block_height, node.signature_key)
351                    .await
352                {
353                    drb_leaf = Some(leaf);
354                    break;
355                }
356            }
357
358            match drb_leaf {
359                Some(leaf) => Ok(leaf.next_drb_result.expect(
360                    "We fetched a leaf that is missing a DRB result. This should be impossible.",
361                )),
362                None => {
363                    anyhow::bail!(
364                        "Failed to fetch leaf from all nodes. Height: {drb_block_height}"
365                    );
366                },
367            }
368        }
369    }
370
371    fn add_da_committee(&mut self, first_epoch: u64, committee: Vec<PeerConfig<TYPES>>) {
372        self.inner
373            .add_da_committee(first_epoch, committee.into_iter().map(Into::into).collect());
374    }
375}
hotshot_example_types/membership/strict_membership.rs

hotshot_example_types/membership/
strict_membership.rs
