|
| 1 | +use crate::throughput::Throughput; |
| 2 | +use alto_types::{Block, Scheme}; |
| 3 | +use commonware_consensus::{ |
| 4 | + marshal::{self, Update}, |
| 5 | + types::Height, |
| 6 | + Reporter, |
| 7 | +}; |
| 8 | +use commonware_runtime::{spawn_cell, Clock, ContextCell, Handle, Spawner}; |
| 9 | +use commonware_utils::Acknowledgement; |
| 10 | +use futures::{channel::mpsc, SinkExt, StreamExt}; |
| 11 | +use tracing::info; |
| 12 | + |
| 13 | +const THROUGHPUT_WINDOW: std::time::Duration = std::time::Duration::from_secs(30); |
| 14 | +const PRUNE_INTERVAL: u64 = 10_000; |
| 15 | + |
| 16 | +/// Formats an estimated time of arrival (ETA) based on the remaining work and rate. |
| 17 | +fn format_eta(remaining: u64, rate: f64) -> String { |
| 18 | + if remaining == 0 { |
| 19 | + return "0s".to_string(); |
| 20 | + } |
| 21 | + if !rate.is_finite() || rate <= 0.0 { |
| 22 | + return "unknown".to_string(); |
| 23 | + } |
| 24 | + |
| 25 | + let secs = (remaining as f64 / rate) as u64; |
| 26 | + let (h, m, s) = (secs / 3600, (secs % 3600) / 60, secs % 60); |
| 27 | + if h > 0 { |
| 28 | + format!("{h}h{m:02}m{s:02}s") |
| 29 | + } else if m > 0 { |
| 30 | + format!("{m}m{s:02}s") |
| 31 | + } else { |
| 32 | + format!("{s}s") |
| 33 | + } |
| 34 | +} |
| 35 | + |
| 36 | +/// Formats ETA when remaining work may be unknown (e.g. tip not received yet). |
| 37 | +fn format_eta_maybe(remaining: Option<u64>, rate: f64) -> String { |
| 38 | + match remaining { |
| 39 | + Some(remaining) => format_eta(remaining, rate), |
| 40 | + None => "unknown".to_string(), |
| 41 | + } |
| 42 | +} |
| 43 | + |
| 44 | +/// A forwarder of [Update] messages to the [Application]. |
| 45 | +#[derive(Clone)] |
| 46 | +pub(crate) struct Mailbox { |
| 47 | + tx: mpsc::Sender<Update<Block>>, |
| 48 | +} |
| 49 | + |
| 50 | +impl Reporter for Mailbox { |
| 51 | + type Activity = Update<Block>; |
| 52 | + |
| 53 | + async fn report(&mut self, activity: Self::Activity) { |
| 54 | + let _ = self.tx.send(activity).await; |
| 55 | + } |
| 56 | +} |
| 57 | + |
| 58 | +/// A simple application that tracks just tracks the rate of block processing. |
| 59 | +pub(crate) struct Application<E: Clock + Spawner> { |
| 60 | + context: ContextCell<E>, |
| 61 | + rx: mpsc::Receiver<Update<Block>>, |
| 62 | + throughput: Throughput, |
| 63 | + tip: Option<Height>, |
| 64 | + mailbox: marshal::Mailbox<Scheme, Block>, |
| 65 | + pruning_depth: Option<u64>, |
| 66 | +} |
| 67 | + |
| 68 | +impl<E: Clock + Spawner> Application<E> { |
| 69 | + pub(crate) fn new( |
| 70 | + context: E, |
| 71 | + mailbox: marshal::Mailbox<Scheme, Block>, |
| 72 | + mailbox_size: usize, |
| 73 | + pruning_depth: Option<u64>, |
| 74 | + ) -> (Self, Mailbox) { |
| 75 | + let (tx, rx) = mpsc::channel(mailbox_size); |
| 76 | + let app = Self { |
| 77 | + context: ContextCell::new(context.clone()), |
| 78 | + rx, |
| 79 | + throughput: Throughput::new(THROUGHPUT_WINDOW), |
| 80 | + tip: None, |
| 81 | + mailbox, |
| 82 | + pruning_depth, |
| 83 | + }; |
| 84 | + (app, Mailbox { tx }) |
| 85 | + } |
| 86 | + |
| 87 | + pub(crate) fn start(mut self) -> Handle<()> { |
| 88 | + spawn_cell!(self.context, self.run().await) |
| 89 | + } |
| 90 | + |
| 91 | + async fn run(mut self) { |
| 92 | + while let Some(msg) = self.rx.next().await { |
| 93 | + match msg { |
| 94 | + Update::Tip(_, height, _) => { |
| 95 | + self.tip = Some(height); |
| 96 | + } |
| 97 | + Update::Block(block, ack) => { |
| 98 | + // This is where an application would process the |
| 99 | + // finalized block (e.g. update state, index transactions, |
| 100 | + // serve queries, etc.). |
| 101 | + let height = block.height.get(); |
| 102 | + let bps = self.throughput.record(self.context.current()); |
| 103 | + let remaining = self.tip.map(|t| t.get().saturating_sub(height)); |
| 104 | + info!( |
| 105 | + height, |
| 106 | + tip = self.tip.map(|h| h.get()), |
| 107 | + bps = %format_args!("{bps:.2}"), |
| 108 | + eta = %format_args!("{}", format_eta_maybe(remaining, bps)), |
| 109 | + "processed block" |
| 110 | + ); |
| 111 | + ack.acknowledge(); |
| 112 | + |
| 113 | + // Prune the archive if the height is a multiple of the prune interval. |
| 114 | + if let Some(depth) = self.pruning_depth.filter(|_| height % PRUNE_INTERVAL == 0) |
| 115 | + { |
| 116 | + let prune_to = height.saturating_sub(depth); |
| 117 | + if prune_to > 0 { |
| 118 | + self.mailbox.prune(Height::new(prune_to)).await; |
| 119 | + } |
| 120 | + } |
| 121 | + } |
| 122 | + } |
| 123 | + } |
| 124 | + } |
| 125 | +} |
| 126 | + |
| 127 | +#[cfg(test)] |
| 128 | +mod tests { |
| 129 | + use super::{format_eta, format_eta_maybe}; |
| 130 | + |
| 131 | + #[test] |
| 132 | + fn eta_is_unknown_when_rate_is_zero_and_remaining_non_zero() { |
| 133 | + assert_eq!(format_eta(42, 0.0), "unknown"); |
| 134 | + } |
| 135 | + |
| 136 | + #[test] |
| 137 | + fn eta_is_zero_when_no_remaining_work() { |
| 138 | + assert_eq!(format_eta(0, 0.0), "0s"); |
| 139 | + } |
| 140 | + |
| 141 | + #[test] |
| 142 | + fn eta_is_unknown_when_remaining_is_unknown() { |
| 143 | + assert_eq!(format_eta_maybe(None, 123.0), "unknown"); |
| 144 | + } |
| 145 | +} |
0 commit comments