Skip to main content

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