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#[async_trait]
41pub trait NewProtocolStorage<T: NodeType>: StorageTrait<T> {
42 async fn append_high_qc2(&self, high_qc: Certificate1<T>) -> anyhow::Result<()>;
45
46 async fn load_high_qc2(&self) -> anyhow::Result<Option<Certificate1<T>>>;
48
49 async fn append_boundary_qc2(&self, qc: Certificate1<T>) -> anyhow::Result<()>;
57
58 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 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 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
99pub 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 pub fn append_boundary_qc2(&mut self, qc: Certificate1<T>) {
239 let view = qc.view_number();
240 let storage = self.storage.clone();
241 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 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 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 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 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}