hotshot_new_protocol/
storage.rs

1use std::{collections::BTreeMap, marker::PhantomData, sync::Arc, time::Duration};
2
3use async_trait::async_trait;
4use committable::Commitment;
5use hotshot::{traits::BlockPayload, types::SignatureKey};
6use hotshot_example_types::storage_types::TestStorage;
7use hotshot_types::{
8    data::{
9        DaProposal2, EpochNumber, Leaf2, QuorumProposal2, QuorumProposalWrapper, VidCommitment,
10        VidDisperseShare, VidDisperseShare2, ViewChangeEvidence2, ViewNumber,
11    },
12    event::HotShotAction,
13    message::Proposal as SignedProposal,
14    simple_certificate::LightClientStateUpdateCertificateV2,
15    traits::{
16        EncodeBytes,
17        metrics::{Histogram, Metrics},
18        node_implementation::NodeType,
19        storage::Storage as StorageTrait,
20    },
21    utils::EpochTransitionIndicator,
22    vote::HasViewNumber,
23};
24use tokio::{
25    task::{AbortHandle, JoinSet},
26    time::sleep,
27};
28use tracing::{error, info, warn};
29
30use crate::{
31    coordinator::metrics::{Measurement, finish_measurement},
32    helpers::proposal_commitment,
33    message::{Certificate1, Certificate2, Proposal},
34};
35
36const RETRY_DELAY: Duration = Duration::from_millis(300);
37
38/// New protocol storage extension for data that is not part of the legacy HotShot storage trait.
39#[async_trait]
40pub trait NewProtocolStorage<T: NodeType>: StorageTrait<T> {
41    async fn append_cert2(&self, view: ViewNumber, cert: Certificate2<T>) -> anyhow::Result<()>;
42
43    /// Persist the locked QC, written before each phase-2 vote so the lock
44    /// survives a restart instead of regressing to the decided-anchor QC.
45    async fn append_high_qc2(&self, high_qc: Certificate1<T>) -> anyhow::Result<()>;
46
47    /// Load the persisted locked QC, if any.
48    async fn load_high_qc2(&self) -> anyhow::Result<Option<Certificate1<T>>>;
49}
50
51#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
52pub enum ActionKind {
53    Vote,
54    Propose,
55}
56
57impl From<ActionKind> for HotShotAction {
58    fn from(kind: ActionKind) -> Self {
59        match kind {
60            ActionKind::Vote => HotShotAction::Vote,
61            ActionKind::Propose => HotShotAction::Propose,
62        }
63    }
64}
65
66#[derive(Clone, Debug, PartialEq, Eq)]
67pub enum StorageOutput<T: NodeType> {
68    Proposal(ViewNumber, Commitment<Leaf2<T>>),
69    Vid(ViewNumber),
70    Action(ViewNumber, ActionKind),
71    /// The locked QC for the given view has been durably persisted.
72    HighQc(ViewNumber),
73}
74
75impl<T: NodeType> StorageOutput<T> {
76    pub fn view_number(&self) -> ViewNumber {
77        match self {
78            Self::Proposal(view, _)
79            | Self::Vid(view)
80            | Self::Action(view, _)
81            | Self::HighQc(view) => *view,
82        }
83    }
84}
85
86/// Request-to-completion latency of each storage operation, including task
87/// queueing, backend lock waits, and error-retry loops. These latencies gate
88/// consensus progress: proposal release waits on `append_proposal` and
89/// `record_action(Propose)`, and phase-2 votes wait on `append_vid`,
90/// `record_action(Vote)`, and `append_high_qc2`.
91///
92/// [`Measurement`] records on drop, so an op aborted by GC or shutdown still
93/// records its elapsed time up to the abort.
94pub struct StorageMetrics {
95    append_vid: Arc<dyn Histogram>,
96    append_da: Arc<dyn Histogram>,
97    append_cert2: Arc<dyn Histogram>,
98    append_high_qc: Arc<dyn Histogram>,
99    append_state_cert: Arc<dyn Histogram>,
100    append_proposal: Arc<dyn Histogram>,
101    record_action: Arc<dyn Histogram>,
102}
103
104impl StorageMetrics {
105    pub fn new(m: &dyn Metrics) -> Self {
106        let histogram = |name: &str| -> Arc<dyn Histogram> {
107            m.create_histogram(format!("storage_{name}"), Some("s".into()))
108                .into()
109        };
110        Self {
111            append_vid: histogram("append_vid"),
112            append_da: histogram("append_da"),
113            append_cert2: histogram("append_cert2"),
114            append_high_qc: histogram("append_high_qc"),
115            append_state_cert: histogram("append_state_cert"),
116            append_proposal: histogram("append_proposal"),
117            record_action: histogram("record_action"),
118        }
119    }
120}
121
122pub struct Storage<T: NodeType, S> {
123    storage: S,
124    private_key: <T::SignatureKey as SignatureKey>::PrivateKey,
125    tasks: JoinSet<Option<StorageOutput<T>>>,
126    handles: BTreeMap<ViewNumber, Vec<AbortHandle>>,
127    metrics: Option<StorageMetrics>,
128}
129
130impl<T: NodeType, S: NewProtocolStorage<T>> Storage<T, S> {
131    pub fn new(storage: S, private_key: <T::SignatureKey as SignatureKey>::PrivateKey) -> Self {
132        Self {
133            storage,
134            private_key,
135            tasks: JoinSet::new(),
136            handles: BTreeMap::new(),
137            metrics: None,
138        }
139    }
140
141    pub fn with_metrics(mut self, m: &dyn Metrics) -> Self {
142        if m.is_recording() {
143            self.metrics = Some(StorageMetrics::new(m));
144        }
145        self
146    }
147
148    pub fn append_vid(&mut self, vid_share: VidDisperseShare2<T>) {
149        let view = vid_share.view_number;
150        let storage = self.storage.clone();
151        let private_key = self.private_key.clone();
152        let timer = self
153            .metrics
154            .as_ref()
155            .map(|m| Measurement::start(m.append_vid.clone()));
156        let handle = self.tasks.spawn(async move {
157            let share: VidDisperseShare<T> = VidDisperseShare::V2(vid_share);
158            let Some(proposal) = share.to_proposal(&private_key) else {
159                error!("failed to sign VID share for storage");
160                return None;
161            };
162            loop {
163                match storage.append_vid(&proposal).await {
164                    Ok(()) => {
165                        finish_measurement(timer);
166                        return Some(StorageOutput::Vid(view));
167                    },
168                    Err(err) => {
169                        warn!(%err, "failed to append VID share, retrying");
170                        sleep(RETRY_DELAY).await;
171                    },
172                }
173            }
174        });
175        self.handles.entry(view).or_default().push(handle);
176    }
177
178    pub fn append_da(
179        &mut self,
180        view_number: ViewNumber,
181        epoch: EpochNumber,
182        block_payload: T::BlockPayload,
183        metadata: <T::BlockPayload as BlockPayload<T>>::Metadata,
184        vid_commit: VidCommitment,
185    ) {
186        let storage = self.storage.clone();
187        let private_key = self.private_key.clone();
188        let timer = self
189            .metrics
190            .as_ref()
191            .map(|m| Measurement::start(m.append_da.clone()));
192        let handle = self.tasks.spawn(async move {
193            let data = DaProposal2 {
194                encoded_transactions: block_payload.encode(),
195                metadata,
196                view_number,
197                epoch: Some(epoch),
198                epoch_transition_indicator: EpochTransitionIndicator::NotInTransition,
199            };
200            let Ok(signature) = T::SignatureKey::sign(&private_key, &[]) else {
201                error!("failed to sign DA proposal for storage");
202                return None;
203            };
204            let proposal = SignedProposal {
205                data,
206                signature,
207                _pd: PhantomData,
208            };
209            loop {
210                match storage.append_da2(&proposal, vid_commit).await {
211                    Ok(()) => {
212                        finish_measurement(timer);
213                        return None;
214                    },
215                    Err(err) => {
216                        warn!(%err, "failed to append DA proposal, retrying");
217                        sleep(RETRY_DELAY).await;
218                    },
219                }
220            }
221        });
222        self.handles.entry(view_number).or_default().push(handle);
223    }
224
225    pub fn append_cert2(&mut self, view: ViewNumber, cert2: Certificate2<T>) {
226        let storage = self.storage.clone();
227        let timer = self
228            .metrics
229            .as_ref()
230            .map(|m| Measurement::start(m.append_cert2.clone()));
231        let handle = self.tasks.spawn(async move {
232            loop {
233                match storage.append_cert2(view, cert2.clone()).await {
234                    Ok(()) => {
235                        finish_measurement(timer);
236                        return None;
237                    },
238                    Err(err) => {
239                        warn!(%err, %view, "failed to append cert2, retrying");
240                        sleep(RETRY_DELAY).await;
241                    },
242                }
243            }
244        });
245        self.handles.entry(view).or_default().push(handle);
246    }
247
248    /// Persist the locked QC; on success emits [`StorageOutput::HighQc`], which
249    /// gates the matching phase-2 vote.
250    pub fn append_high_qc2(&mut self, high_qc: Certificate1<T>) {
251        let view = high_qc.view_number();
252        let storage = self.storage.clone();
253        let timer = self
254            .metrics
255            .as_ref()
256            .map(|m| Measurement::start(m.append_high_qc.clone()));
257        let handle = self.tasks.spawn(async move {
258            loop {
259                match storage.append_high_qc2(high_qc.clone()).await {
260                    Ok(()) => {
261                        finish_measurement(timer);
262                        return Some(StorageOutput::HighQc(view));
263                    },
264                    Err(err) => {
265                        warn!(%err, %view, "failed to append high qc, retrying");
266                        sleep(RETRY_DELAY).await;
267                    },
268                }
269            }
270        });
271        self.handles.entry(view).or_default().push(handle);
272    }
273
274    pub fn append_state_cert(
275        &mut self,
276        view: ViewNumber,
277        state_cert: LightClientStateUpdateCertificateV2<T>,
278    ) {
279        let storage = self.storage.clone();
280        let timer = self
281            .metrics
282            .as_ref()
283            .map(|m| Measurement::start(m.append_state_cert.clone()));
284        let handle = self.tasks.spawn(async move {
285            loop {
286                match storage.update_state_cert(state_cert.clone()).await {
287                    Ok(()) => {
288                        finish_measurement(timer);
289                        return None;
290                    },
291                    Err(err) => {
292                        warn!(%err, epoch = %state_cert.epoch, "failed to append state cert, retrying");
293                        sleep(RETRY_DELAY).await;
294                    },
295                }
296            }
297        });
298        self.handles.entry(view).or_default().push(handle);
299    }
300
301    pub fn append_proposal(&mut self, proposal: Proposal<T>) {
302        let view = proposal.view_number;
303        let commitment = proposal_commitment(&proposal);
304        let storage = self.storage.clone();
305        let private_key = self.private_key.clone();
306        let timer = self
307            .metrics
308            .as_ref()
309            .map(|m| Measurement::start(m.append_proposal.clone()));
310        let handle = self.tasks.spawn(async move {
311            let data = QuorumProposalWrapper {
312                proposal: QuorumProposal2 {
313                    block_header: proposal.block_header,
314                    view_number: proposal.view_number,
315                    epoch: Some(proposal.epoch),
316                    justify_qc: proposal.justify_qc,
317                    next_epoch_justify_qc: None,
318                    upgrade_certificate: proposal.upgrade_certificate,
319                    view_change_evidence: proposal
320                        .view_change_evidence
321                        .map(ViewChangeEvidence2::Timeout),
322                    next_drb_result: proposal.next_drb_result,
323                    state_cert: proposal.state_cert,
324                },
325            };
326            let Ok(signature) = T::SignatureKey::sign(&private_key, &[]) else {
327                error!("failed to sign quorum proposal for storage");
328                return None;
329            };
330            let signed = SignedProposal {
331                data,
332                signature,
333                _pd: PhantomData,
334            };
335            loop {
336                match storage.append_proposal_wrapper(&signed).await {
337                    Ok(()) => {
338                        finish_measurement(timer);
339                        return Some(StorageOutput::Proposal(view, commitment));
340                    },
341                    Err(err) => {
342                        warn!(%err, "failed to append proposal, retrying");
343                        sleep(RETRY_DELAY).await;
344                    },
345                }
346            }
347        });
348        self.handles.entry(view).or_default().push(handle);
349    }
350
351    pub fn record_action(
352        &mut self,
353        view: ViewNumber,
354        epoch: Option<EpochNumber>,
355        kind: ActionKind,
356    ) {
357        let storage = self.storage.clone();
358        let timer = self
359            .metrics
360            .as_ref()
361            .map(|m| Measurement::start(m.record_action.clone()));
362        let handle = self.tasks.spawn(async move {
363            loop {
364                match storage.record_action(view, epoch, kind.into()).await {
365                    Ok(()) => {
366                        finish_measurement(timer);
367                        return Some(StorageOutput::Action(view, kind));
368                    },
369                    Err(err) => {
370                        warn!(%err, %view, ?kind, "failed to record action, retrying");
371                        sleep(RETRY_DELAY).await;
372                    },
373                }
374            }
375        });
376        self.handles.entry(view).or_default().push(handle);
377    }
378
379    pub async fn next(&mut self) -> Option<StorageOutput<T>> {
380        loop {
381            match self.tasks.join_next().await? {
382                Ok(Some(output)) => return Some(output),
383                Ok(None) | Err(_) => continue,
384            }
385        }
386    }
387
388    pub fn gc(&mut self, view_number: ViewNumber) {
389        let keep = self.handles.split_off(&view_number);
390        for handles in self.handles.values() {
391            for handle in handles {
392                handle.abort();
393            }
394        }
395        self.handles = keep;
396    }
397
398    /// Wait for all current storage writes to complete.
399    pub async fn flush(mut self) {
400        info!(
401            tasks = self.tasks.len(),
402            "flushing storage tasks during shutdown"
403        );
404        while let Some(result) = self.tasks.join_next().await {
405            if let Err(err) = result {
406                warn!(%err, "storage task failed during shutdown");
407            }
408        }
409        info!("storage flush complete");
410    }
411}
412
413#[async_trait]
414impl<T: NodeType> NewProtocolStorage<T> for TestStorage<T> {
415    async fn append_cert2(&self, _view: ViewNumber, _cert: Certificate2<T>) -> anyhow::Result<()> {
416        Ok(())
417    }
418
419    async fn append_high_qc2(&self, high_qc: Certificate1<T>) -> anyhow::Result<()> {
420        // `Certificate1<T>` is a `QuorumCertificate2<T>`; reuse the monotonic legacy slot.
421        StorageTrait::update_high_qc2(self, high_qc).await
422    }
423
424    async fn load_high_qc2(&self) -> anyhow::Result<Option<Certificate1<T>>> {
425        Ok(self.high_qc_cloned().await)
426    }
427}
hotshot_new_protocol/storage.rs

hotshot_new_protocol/
storage.rs
