hotshot_example_types/
storage_types.rs

1// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
2// This file is part of the HotShot repository.
3
4// You should have received a copy of the MIT License
5// along with the HotShot repository. If not, see <https://mit-license.org/>.
6
7use std::{
8    collections::{BTreeMap, HashMap},
9    sync::{
10        Arc,
11        atomic::{AtomicBool, Ordering},
12    },
13};
14
15use anyhow::{Result, anyhow, bail};
16use async_lock::RwLock;
17use async_trait::async_trait;
18use hotshot_types::{
19    data::{
20        DaProposal, DaProposal2, EpochNumber, Leaf2, QuorumProposal, QuorumProposal2,
21        QuorumProposalWrapper, VidCommitment, VidDisperseShare, ViewNumber,
22    },
23    drb::{DrbInput, DrbResult},
24    event::HotShotAction,
25    message::{Proposal, convert_proposal},
26    simple_certificate::{
27        LightClientStateUpdateCertificateV2, NextEpochQuorumCertificate2, QuorumCertificate2,
28        UpgradeCertificate,
29    },
30    traits::{node_implementation::NodeType, storage::Storage},
31    vote::HasViewNumber,
32};
33
34use crate::testable_delay::{DelayConfig, SupportedTraitTypesForAsyncDelay, TestableDelay};
35
36type VidShares<TYPES> = BTreeMap<
37    ViewNumber,
38    HashMap<<TYPES as NodeType>::SignatureKey, Proposal<TYPES, VidDisperseShare<TYPES>>>,
39>;
40#[derive(Clone, Debug)]
41pub struct TestStorageState<TYPES: NodeType> {
42    vids: VidShares<TYPES>,
43    das: HashMap<ViewNumber, Proposal<TYPES, DaProposal<TYPES>>>,
44    da2s: HashMap<ViewNumber, Proposal<TYPES, DaProposal2<TYPES>>>,
45    pub proposals: BTreeMap<ViewNumber, Proposal<TYPES, QuorumProposal<TYPES>>>,
46    pub proposals2: BTreeMap<ViewNumber, Proposal<TYPES, QuorumProposal2<TYPES>>>,
47    pub proposals_wrapper: BTreeMap<ViewNumber, Proposal<TYPES, QuorumProposalWrapper<TYPES>>>,
48    high_qc: Option<hotshot_types::simple_certificate::QuorumCertificate<TYPES>>,
49    high_qc2: Option<hotshot_types::simple_certificate::QuorumCertificate2<TYPES>>,
50    eqc: Option<(
51        hotshot_types::simple_certificate::QuorumCertificate2<TYPES>,
52        hotshot_types::simple_certificate::NextEpochQuorumCertificate2<TYPES>,
53    )>,
54    next_epoch_high_qc2:
55        Option<hotshot_types::simple_certificate::NextEpochQuorumCertificate2<TYPES>>,
56    action: ViewNumber,
57    action_log: Vec<(ViewNumber, HotShotAction)>,
58    epoch: Option<EpochNumber>,
59    state_certs: BTreeMap<EpochNumber, LightClientStateUpdateCertificateV2<TYPES>>,
60    drb_results: BTreeMap<EpochNumber, DrbResult>,
61    drb_inputs: BTreeMap<u64, DrbInput>,
62    epoch_roots: BTreeMap<EpochNumber, TYPES::BlockHeader>,
63    restart_view: ViewNumber,
64    anchor_leaf: Option<(Leaf2<TYPES>, QuorumCertificate2<TYPES>)>,
65}
66
67impl<TYPES: NodeType> Default for TestStorageState<TYPES> {
68    fn default() -> Self {
69        Self {
70            vids: BTreeMap::new(),
71            das: HashMap::new(),
72            da2s: HashMap::new(),
73            proposals: BTreeMap::new(),
74            proposals2: BTreeMap::new(),
75            proposals_wrapper: BTreeMap::new(),
76            high_qc: None,
77            high_qc2: None,
78            eqc: None,
79            next_epoch_high_qc2: None,
80            action: ViewNumber::genesis(),
81            action_log: Vec::new(),
82            epoch: None,
83            state_certs: BTreeMap::new(),
84            drb_results: BTreeMap::new(),
85            drb_inputs: BTreeMap::new(),
86            epoch_roots: BTreeMap::new(),
87            restart_view: ViewNumber::genesis(),
88            anchor_leaf: None,
89        }
90    }
91}
92
93#[derive(Clone, Debug)]
94pub struct TestStorage<TYPES: NodeType> {
95    pub inner: Arc<RwLock<TestStorageState<TYPES>>>,
96    /// `should_return_err` is a testing utility to validate negative cases.
97    pub should_return_err: Arc<AtomicBool>,
98    pub delay_config: DelayConfig,
99    pub decided_upgrade_certificate: Arc<RwLock<Option<UpgradeCertificate<TYPES>>>>,
100}
101
102impl<TYPES: NodeType> Default for TestStorage<TYPES> {
103    fn default() -> Self {
104        Self {
105            inner: Arc::new(RwLock::new(TestStorageState::default())),
106            should_return_err: Arc::new(AtomicBool::new(false)),
107            delay_config: DelayConfig::default(),
108            decided_upgrade_certificate: Arc::new(RwLock::new(None)),
109        }
110    }
111}
112
113#[async_trait]
114impl<TYPES: NodeType> TestableDelay for TestStorage<TYPES> {
115    async fn run_delay_settings_from_config(delay_config: &DelayConfig) {
116        if let Some(settings) = delay_config.get_setting(&SupportedTraitTypesForAsyncDelay::Storage)
117        {
118            Self::handle_async_delay(settings).await;
119        }
120    }
121}
122
123impl<TYPES: NodeType> TestStorage<TYPES> {
124    pub async fn proposals_cloned(
125        &self,
126    ) -> BTreeMap<ViewNumber, Proposal<TYPES, QuorumProposalWrapper<TYPES>>> {
127        self.inner.read().await.proposals_wrapper.clone()
128    }
129
130    pub async fn high_qc_cloned(&self) -> Option<QuorumCertificate2<TYPES>> {
131        self.inner.read().await.high_qc2.clone()
132    }
133
134    pub async fn next_epoch_high_qc_cloned(&self) -> Option<NextEpochQuorumCertificate2<TYPES>> {
135        self.inner.read().await.next_epoch_high_qc2.clone()
136    }
137
138    pub async fn decided_upgrade_certificate(&self) -> Option<UpgradeCertificate<TYPES>> {
139        self.decided_upgrade_certificate.read().await.clone()
140    }
141
142    pub async fn last_actioned_view(&self) -> ViewNumber {
143        self.inner.read().await.action
144    }
145
146    /// Every action successfully recorded, in call order.
147    pub async fn action_log(&self) -> Vec<(ViewNumber, HotShotAction)> {
148        self.inner.read().await.action_log.clone()
149    }
150
151    pub async fn restart_view(&self) -> ViewNumber {
152        self.inner.read().await.restart_view
153    }
154
155    /// Record the newest decided leaf and the QC certifying it, keeping the
156    /// pair with the highest view. Plays the role of the application-level
157    /// anchor persistence (e.g. `append_decided_leaves` in the sequencer);
158    /// the anchor is what a restarted node resumes consensus from.
159    pub async fn update_anchor_leaf(&self, leaf: Leaf2<TYPES>, qc: QuorumCertificate2<TYPES>) {
160        let mut inner = self.inner.write().await;
161        if inner
162            .anchor_leaf
163            .as_ref()
164            .is_none_or(|(anchor, _)| leaf.view_number() > anchor.view_number())
165        {
166            inner.anchor_leaf = Some((leaf, qc));
167        }
168    }
169
170    pub async fn anchor_leaf(&self) -> Option<(Leaf2<TYPES>, QuorumCertificate2<TYPES>)> {
171        self.inner.read().await.anchor_leaf.clone()
172    }
173
174    pub async fn last_actioned_epoch(&self) -> Option<EpochNumber> {
175        self.inner.read().await.epoch
176    }
177    pub async fn vids_cloned(&self) -> VidShares<TYPES> {
178        self.inner.read().await.vids.clone()
179    }
180
181    pub async fn state_cert_cloned(&self) -> Option<LightClientStateUpdateCertificateV2<TYPES>> {
182        self.inner
183            .read()
184            .await
185            .state_certs
186            .iter()
187            .next_back()
188            .map(|(_, cert)| cert.clone())
189    }
190}
191
192#[async_trait]
193impl<TYPES: NodeType> Storage<TYPES> for TestStorage<TYPES> {
194    async fn append_vid(&self, proposal: &Proposal<TYPES, VidDisperseShare<TYPES>>) -> Result<()> {
195        if self.should_return_err.load(Ordering::Relaxed) {
196            bail!("Failed to append VID proposal to storage");
197        }
198        Self::run_delay_settings_from_config(&self.delay_config).await;
199        let mut inner = self.inner.write().await;
200        inner
201            .vids
202            .entry(proposal.data.view_number())
203            .or_default()
204            .insert(proposal.data.recipient_key().clone(), proposal.clone());
205        Ok(())
206    }
207
208    async fn append_da(
209        &self,
210        proposal: &Proposal<TYPES, DaProposal<TYPES>>,
211        _vid_commit: VidCommitment,
212    ) -> Result<()> {
213        if self.should_return_err.load(Ordering::Relaxed) {
214            bail!("Failed to append DA proposal to storage");
215        }
216        Self::run_delay_settings_from_config(&self.delay_config).await;
217        let mut inner = self.inner.write().await;
218        inner
219            .das
220            .insert(proposal.data.view_number, proposal.clone());
221        Ok(())
222    }
223
224    async fn append_da2(
225        &self,
226        proposal: &Proposal<TYPES, DaProposal2<TYPES>>,
227        _vid_commit: VidCommitment,
228    ) -> Result<()> {
229        if self.should_return_err.load(Ordering::Relaxed) {
230            bail!("Failed to append DA proposal (2) to storage");
231        }
232        Self::run_delay_settings_from_config(&self.delay_config).await;
233        let mut inner = self.inner.write().await;
234        inner
235            .da2s
236            .insert(proposal.data.view_number, proposal.clone());
237        Ok(())
238    }
239
240    async fn append_proposal(
241        &self,
242        proposal: &Proposal<TYPES, QuorumProposal<TYPES>>,
243    ) -> Result<()> {
244        if self.should_return_err.load(Ordering::Relaxed) {
245            bail!("Failed to append Quorum proposal (1) to storage");
246        }
247        Self::run_delay_settings_from_config(&self.delay_config).await;
248        let mut inner = self.inner.write().await;
249        inner
250            .proposals
251            .insert(proposal.data.view_number, proposal.clone());
252        Ok(())
253    }
254
255    async fn append_proposal2(
256        &self,
257        proposal: &Proposal<TYPES, QuorumProposal2<TYPES>>,
258    ) -> Result<()> {
259        if self.should_return_err.load(Ordering::Relaxed) {
260            bail!("Failed to append Quorum proposal (2) to storage");
261        }
262        Self::run_delay_settings_from_config(&self.delay_config).await;
263        let mut inner = self.inner.write().await;
264        inner
265            .proposals2
266            .insert(proposal.data.view_number, proposal.clone());
267        Ok(())
268    }
269
270    async fn append_proposal_wrapper(
271        &self,
272        proposal: &Proposal<TYPES, QuorumProposalWrapper<TYPES>>,
273    ) -> Result<()> {
274        if self.should_return_err.load(Ordering::Relaxed) {
275            bail!("Failed to append Quorum proposal (wrapped) to storage");
276        }
277        Self::run_delay_settings_from_config(&self.delay_config).await;
278        let mut inner = self.inner.write().await;
279        inner
280            .proposals_wrapper
281            .insert(proposal.data.view_number(), proposal.clone());
282        Ok(())
283    }
284
285    async fn record_action(
286        &self,
287        view: ViewNumber,
288        epoch: Option<EpochNumber>,
289        action: hotshot_types::event::HotShotAction,
290    ) -> Result<()> {
291        if self.should_return_err.load(Ordering::Relaxed) {
292            bail!("Failed to append Action to storage");
293        }
294        let mut inner = self.inner.write().await;
295        inner.action_log.push((view, action));
296        if matches!(
297            action,
298            HotShotAction::Vote | HotShotAction::Propose | HotShotAction::TimeoutVote
299        ) {
300            if view > inner.action {
301                inner.action = view;
302            }
303            if epoch > inner.epoch {
304                inner.epoch = epoch;
305            }
306        }
307        if matches!(action, HotShotAction::Vote) {
308            inner.restart_view = view + 1;
309        }
310        Self::run_delay_settings_from_config(&self.delay_config).await;
311        Ok(())
312    }
313
314    async fn update_high_qc(
315        &self,
316        new_high_qc: hotshot_types::simple_certificate::QuorumCertificate<TYPES>,
317    ) -> Result<()> {
318        if self.should_return_err.load(Ordering::Relaxed) {
319            bail!("Failed to update high qc to storage");
320        }
321        Self::run_delay_settings_from_config(&self.delay_config).await;
322        let mut inner = self.inner.write().await;
323        if let Some(ref current_high_qc) = inner.high_qc {
324            if new_high_qc.view_number() > current_high_qc.view_number() {
325                inner.high_qc = Some(new_high_qc);
326            }
327        } else {
328            inner.high_qc = Some(new_high_qc);
329        }
330        Ok(())
331    }
332
333    /// Update the current high QC in storage.
334    async fn update_eqc(
335        &self,
336        high_qc: QuorumCertificate2<TYPES>,
337        next_epoch_high_qc: NextEpochQuorumCertificate2<TYPES>,
338    ) -> Result<()> {
339        if self.should_return_err.load(Ordering::Relaxed) {
340            bail!("Failed to update eqc in storage");
341        }
342        Self::run_delay_settings_from_config(&self.delay_config).await;
343        let mut inner = self.inner.write().await;
344        if let Some((ref current_high_qc, _)) = inner.eqc {
345            if high_qc.view_number() > current_high_qc.view_number() {
346                inner.eqc = Some((high_qc, next_epoch_high_qc));
347            }
348        } else {
349            inner.eqc = Some((high_qc, next_epoch_high_qc));
350        }
351        Ok(())
352    }
353
354    async fn update_high_qc2(
355        &self,
356        new_high_qc: hotshot_types::simple_certificate::QuorumCertificate2<TYPES>,
357    ) -> Result<()> {
358        if self.should_return_err.load(Ordering::Relaxed) {
359            bail!("Failed to update high qc to storage");
360        }
361        Self::run_delay_settings_from_config(&self.delay_config).await;
362        let mut inner = self.inner.write().await;
363        if let Some(ref current_high_qc) = inner.high_qc2 {
364            if new_high_qc.view_number() > current_high_qc.view_number() {
365                inner.high_qc2 = Some(new_high_qc);
366            }
367        } else {
368            inner.high_qc2 = Some(new_high_qc);
369        }
370        Ok(())
371    }
372
373    async fn update_state_cert(
374        &self,
375        state_cert: LightClientStateUpdateCertificateV2<TYPES>,
376    ) -> Result<()> {
377        if self.should_return_err.load(Ordering::Relaxed) {
378            bail!("Failed to update state_cert to storage");
379        }
380        Self::run_delay_settings_from_config(&self.delay_config).await;
381        self.inner
382            .write()
383            .await
384            .state_certs
385            .insert(state_cert.epoch, state_cert);
386        Ok(())
387    }
388
389    async fn update_next_epoch_high_qc2(
390        &self,
391        new_next_epoch_high_qc: hotshot_types::simple_certificate::NextEpochQuorumCertificate2<
392            TYPES,
393        >,
394    ) -> Result<()> {
395        if self.should_return_err.load(Ordering::Relaxed) {
396            bail!("Failed to update next epoch high qc to storage");
397        }
398        Self::run_delay_settings_from_config(&self.delay_config).await;
399        let mut inner = self.inner.write().await;
400        if let Some(ref current_next_epoch_high_qc) = inner.next_epoch_high_qc2 {
401            if new_next_epoch_high_qc.view_number() > current_next_epoch_high_qc.view_number() {
402                inner.next_epoch_high_qc2 = Some(new_next_epoch_high_qc);
403            }
404        } else {
405            inner.next_epoch_high_qc2 = Some(new_next_epoch_high_qc);
406        }
407        Ok(())
408    }
409
410    async fn update_decided_upgrade_certificate(
411        &self,
412        decided_upgrade_certificate: Option<UpgradeCertificate<TYPES>>,
413    ) -> Result<()> {
414        *self.decided_upgrade_certificate.write().await = decided_upgrade_certificate;
415
416        Ok(())
417    }
418
419    async fn migrate_storage(&self) -> Result<()> {
420        let mut storage_writer = self.inner.write().await;
421
422        for (view, proposal) in storage_writer.proposals.clone().iter() {
423            storage_writer
424                .proposals2
425                .insert(*view, convert_proposal(proposal.clone()));
426        }
427
428        Ok(())
429    }
430
431    async fn store_drb_result(&self, epoch: EpochNumber, drb_result: DrbResult) -> Result<()> {
432        let mut inner = self.inner.write().await;
433
434        inner.drb_results.insert(epoch, drb_result);
435
436        Ok(())
437    }
438
439    async fn store_epoch_root(
440        &self,
441        epoch: EpochNumber,
442        block_header: TYPES::BlockHeader,
443    ) -> Result<()> {
444        let mut inner = self.inner.write().await;
445
446        inner.epoch_roots.insert(epoch, block_header);
447
448        Ok(())
449    }
450
451    async fn store_drb_input(&self, drb_input: DrbInput) -> Result<()> {
452        let mut inner = self.inner.write().await;
453
454        inner.drb_inputs.insert(drb_input.epoch, drb_input);
455
456        Ok(())
457    }
458
459    async fn load_drb_input(&self, epoch: u64) -> Result<DrbInput> {
460        let inner = self.inner.read().await;
461
462        match inner.drb_inputs.get(&epoch) {
463            Some(drb_input) => Ok(drb_input.clone()),
464            None => Err(anyhow!("Missing DrbInput for epoch {}", epoch)),
465        }
466    }
467}
hotshot_example_types/storage_types.rs

hotshot_example_types/
storage_types.rs
