Skip to main content

hotshot_types/traits/
storage.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
7//! Abstract storage type for storing DA proposals and VID shares
8//!
9//! This modules provides the [`Storage`] trait.
10//!
11
12use std::{sync::Arc, time::Duration};
13
14use anyhow::{Result, anyhow, ensure};
15use async_trait::async_trait;
16use futures::future::BoxFuture;
17use tokio::time::sleep;
18
19use super::node_implementation::NodeType;
20use crate::{
21    data::{
22        DaProposal, DaProposal2, EpochNumber, QuorumProposalWrapper, VidCommitment,
23        VidDisperseShare, ViewNumber,
24    },
25    drb::{DrbInput, DrbResult},
26    event::HotShotAction,
27    message::{Proposal, convert_proposal},
28    simple_certificate::{LightClientStateUpdateCertificateV2, UpgradeCertificate},
29};
30
31/// Abstraction for storing a variety of consensus payload datum.
32#[async_trait]
33pub trait Storage<TYPES: NodeType>: Send + Sync + Clone + 'static {
34    /// Add a proposal to the stored VID proposals.
35    async fn append_vid(&self, proposal: &Proposal<TYPES, VidDisperseShare<TYPES>>) -> Result<()>;
36
37    /// Add a proposal to the stored DA proposals.
38    async fn append_da(
39        &self,
40        proposal: &Proposal<TYPES, DaProposal<TYPES>>,
41        vid_commit: VidCommitment,
42    ) -> Result<()>;
43    /// Add a proposal to the stored DA proposals.
44    async fn append_da2(
45        &self,
46        proposal: &Proposal<TYPES, DaProposal2<TYPES>>,
47        vid_commit: VidCommitment,
48    ) -> Result<()> {
49        self.append_da(&convert_proposal(proposal.clone()), vid_commit)
50            .await
51    }
52    /// Add a proposal we sent to the store
53    async fn append_proposal_wrapper(
54        &self,
55        proposal: &Proposal<TYPES, QuorumProposalWrapper<TYPES>>,
56    ) -> Result<()>;
57    /// Record a HotShotAction taken.
58    async fn record_action(&self, view: ViewNumber, action: HotShotAction) -> Result<()>;
59    /// Update the light client state update certificate in storage.
60    async fn update_state_cert(
61        &self,
62        state_cert: LightClientStateUpdateCertificateV2<TYPES>,
63    ) -> Result<()>;
64    /// Upgrade the current decided upgrade certificate in storage.
65    async fn update_decided_upgrade_certificate(
66        &self,
67        decided_upgrade_certificate: Option<UpgradeCertificate<TYPES>>,
68    ) -> Result<()>;
69    /// Add a drb result
70    async fn store_drb_result(&self, epoch: EpochNumber, drb_result: DrbResult) -> Result<()>;
71    /// Add an epoch block header
72    async fn store_epoch_root(
73        &self,
74        epoch: EpochNumber,
75        block_header: TYPES::BlockHeader,
76    ) -> Result<()>;
77    async fn load_drb_result(&self, epoch: EpochNumber) -> Result<DrbResult> {
78        match self.load_drb_input(*epoch).await {
79            Ok(drb_input) => {
80                ensure!(drb_input.iteration == drb_input.difficulty_level);
81
82                Ok(drb_input.value)
83            },
84            Err(e) => Err(e),
85        }
86    }
87    async fn store_drb_input(&self, drb_input: DrbInput) -> Result<()>;
88    async fn load_drb_input(&self, _epoch: u64) -> Result<DrbInput>;
89}
90
91pub async fn load_drb_input_impl<TYPES: NodeType>(
92    storage: impl Storage<TYPES>,
93    epoch: u64,
94) -> Result<DrbInput> {
95    storage.load_drb_input(epoch).await
96}
97
98pub type LoadDrbProgressFn =
99    std::sync::Arc<dyn Fn(u64) -> BoxFuture<'static, Result<DrbInput>> + Send + Sync>;
100
101pub fn load_drb_progress_fn<TYPES: NodeType>(
102    storage: impl Storage<TYPES> + 'static,
103) -> LoadDrbProgressFn {
104    Arc::new(move |epoch| {
105        let storage = storage.clone();
106        Box::pin(load_drb_input_impl(storage, epoch))
107    })
108}
109
110pub fn null_load_drb_progress_fn() -> LoadDrbProgressFn {
111    Arc::new(move |_drb_input| {
112        Box::pin(async { Err(anyhow!("Using null implementation of load_drb_input")) })
113    })
114}
115
116pub async fn store_drb_input_impl<TYPES: NodeType>(
117    storage: impl Storage<TYPES>,
118    drb_input: DrbInput,
119) -> Result<()> {
120    for attempt in 1..=3 {
121        match storage.store_drb_input(drb_input.clone()).await {
122            Ok(()) => return Ok(()),
123            Err(e) if attempt < 3 => {
124                tracing::warn!("Failed to store DRB input (attempt {attempt}/3): {e}");
125                sleep(Duration::from_millis(300)).await;
126            },
127            Err(e) => {
128                tracing::warn!("Failed to store DRB input (attempt {attempt}/3): {e}");
129                return Err(e);
130            },
131        }
132    }
133
134    Ok(())
135}
136
137pub type StoreDrbProgressFn =
138    std::sync::Arc<dyn Fn(DrbInput) -> BoxFuture<'static, Result<()>> + Send + Sync>;
139
140pub fn store_drb_progress_fn<TYPES: NodeType>(
141    storage: impl Storage<TYPES> + 'static,
142) -> StoreDrbProgressFn {
143    Arc::new(move |drb_input| {
144        let storage = storage.clone();
145        Box::pin(store_drb_input_impl(storage, drb_input))
146    })
147}
148
149pub fn null_store_drb_progress_fn() -> StoreDrbProgressFn {
150    Arc::new(move |_drb_input| Box::pin(async { Ok(()) }))
151}
152
153pub type StoreDrbResultFn = Arc<
154    Box<dyn Fn(EpochNumber, DrbResult) -> BoxFuture<'static, Result<()>> + Send + Sync + 'static>,
155>;
156
157async fn store_drb_result_impl<TYPES: NodeType>(
158    storage: impl Storage<TYPES>,
159    epoch: EpochNumber,
160    drb_result: DrbResult,
161) -> Result<()> {
162    for attempt in 1..=3 {
163        match storage.store_drb_result(epoch, drb_result).await {
164            Ok(()) => return Ok(()),
165            Err(e) if attempt < 3 => {
166                tracing::warn!(
167                    "Failed to store DRB result for epoch {epoch} (attempt {attempt}/3): {e}"
168                );
169                sleep(Duration::from_millis(300)).await;
170            },
171            Err(e) => {
172                tracing::warn!(
173                    "Failed to store DRB result for epoch {epoch} (attempt {attempt}/3): {e}"
174                );
175                return Err(e);
176            },
177        }
178    }
179    Ok(())
180}
181
182/// Helper function to create a callback to add a drb result to storage
183pub fn store_drb_result_fn<TYPES: NodeType>(
184    storage: impl Storage<TYPES> + 'static,
185) -> StoreDrbResultFn {
186    Arc::new(Box::new(move |epoch, drb_result| {
187        let st = storage.clone();
188        Box::pin(store_drb_result_impl(st, epoch, drb_result))
189    }))
190}