more caching

clabby · clabby · commit e87a951cd393 · 2025-10-08T17:24:44.000-04:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/coding/src/lib.rs b/coding/src/lib.rs
@@ -121,7 +121,7 @@ pub trait Scheme: Debug + Clone + Send + Sync + 'static {
     ///
     /// This allows excluding [Scheme::ReShard]s which are invalid, and shouldn't
     /// be considered as progress towards meeting the minimum number of shards.
-    type CheckedShard;
+    type CheckedShard: Clone + Send;
     type Error: std::fmt::Debug;
 
     /// Encode a piece of data, returning a commitment, along with shards, and proofs.
diff --git a/consensus/Cargo.toml b/consensus/Cargo.toml
@@ -30,6 +30,7 @@ commonware-runtime = { workspace = true }
 commonware-storage = { workspace = true, features = ["std"] }
 prometheus-client = { workspace = true }
 governor = { workspace = true }
+rayon = { workspace = true }
 rand = { workspace = true }
 rand_distr = { workspace = true }
 tracing = { workspace = true }
diff --git a/consensus/src/marshal/actor.rs b/consensus/src/marshal/actor.rs
@@ -461,7 +461,17 @@ where
                                 }
                             }
                         }
+                        Message::Notarize { notarization } => {
+                            let commitment = notarization.proposal.payload;
+                            if !shards.has_block(&commitment) {
+                                let start = Instant::now();
+                                let _ = shards.try_reconstruct(commitment).await;
+                                    let elapsed = start.elapsed();
+                                tracing::info!(?elapsed, "Attempted reconstruction");
+                            }
+                        }
                         Message::Notarization { notarization } => {
+                            tracing::warn!(?notarization, "Received notarization");
                             let round = notarization.round();
                             let commitment = notarization.proposal.payload;
 
diff --git a/consensus/src/marshal/ingress/coding/mailbox.rs b/consensus/src/marshal/ingress/coding/mailbox.rs
@@ -65,8 +65,10 @@ where
     /// These blocks are evicted by marshal after they are delivered to the application.
     reconstructed_blocks: BTreeMap<CodingCommitment, CodedBlock<B, S>>,
 
-    block_broadcast_duration: Gauge,
-    shard_broadcast_duration: Gauge,
+    /// Transient caches for progressive block reconstruction.
+    checking_data: BTreeMap<CodingCommitment, S::CheckingData>,
+    checked_shards: BTreeMap<CodingCommitment, BTreeMap<usize, S::CheckedShard>>,
+
     erasure_decode_duration: Gauge,
 }
 
@@ -82,20 +84,6 @@ where
         mailbox: buffered::Mailbox<P, Shard<S, H>>,
         block_codec_cfg: B::Cfg,
     ) -> Self {
-        let block_broadcast_duration = Gauge::default();
-        context.register(
-            "block_broadcast_duration",
-            "Duration of proposer broadcast in milliseconds",
-            block_broadcast_duration.clone(),
-        );
-
-        let shard_broadcast_duration = Gauge::default();
-        context.register(
-            "shard_broadcast_duration",
-            "Duration of individual shard broadcast in milliseconds",
-            shard_broadcast_duration.clone(),
-        );
-
         let erasure_decode_duration = Gauge::default();
         context.register(
             "erasure_decode_duration",
@@ -108,8 +96,8 @@ where
             block_codec_cfg,
             block_subscriptions: BTreeMap::new(),
             reconstructed_blocks: BTreeMap::new(),
-            block_broadcast_duration,
-            shard_broadcast_duration,
+            checking_data: BTreeMap::new(),
+            checked_shards: BTreeMap::new(),
             erasure_decode_duration,
         }
     }
@@ -126,15 +114,12 @@ where
             "number of participants must equal number of shards"
         );
 
-        let start = Instant::now();
         for (index, peer) in participants.into_iter().enumerate() {
             let message = block
                 .shard(index)
                 .expect("peer index impossibly out of bounds");
             let _peers = self.mailbox.broadcast(Recipients::One(peer), message).await;
         }
-        self.block_broadcast_duration
-            .set(start.elapsed().as_millis() as i64);
     }
 
     /// Broadcasts a local [Shard] of a block to all peers, if the [Shard] is present
@@ -172,10 +157,7 @@ where
             // Broadcast the weak shard to all peers for reconstruction.
             let reshard = Shard::new(commitment, index, DistributionShard::Weak(reshard));
 
-            let start = Instant::now();
             let _peers = self.mailbox.broadcast(Recipients::All, reshard).await;
-            self.shard_broadcast_duration
-                .set(start.elapsed().as_millis() as i64);
 
             debug!(%commitment, index, "broadcasted local shard to all peers");
         } else {
@@ -211,24 +193,32 @@ where
         // NOTE: Byzantine peers may send us strong shards as well, but we don't care about those;
         // `Scheme::reshard` verifies the shard against the commitment, and if it doesn't check out,
         // it will be ignored.
-        let Some(checking_data) = shards.iter().find_map(|s| {
-            if let DistributionShard::Strong(shard) = s.deref() {
-                S::reshard(
-                    &config,
-                    &commitment.inner(),
-                    s.index() as u16,
-                    shard.clone(),
-                )
-                .map(|(checking_data, _, _)| checking_data)
-                .ok()
-            } else {
-                None
+        let checking_data = match self.checking_data.entry(commitment) {
+            Entry::Vacant(entry) => {
+                let Some(checking_data) = shards.iter().find_map(|s| {
+                    if let DistributionShard::Strong(shard) = s.deref() {
+                        S::reshard(
+                            &config,
+                            &commitment.inner(),
+                            s.index() as u16,
+                            shard.clone(),
+                        )
+                        .map(|(checking_data, _, _)| checking_data)
+                        .ok()
+                    } else {
+                        None
+                    }
+                }) else {
+                    debug!(%commitment, "No strong shards present to form checking data");
+                    return Ok(None);
+                };
+                entry.insert(checking_data.clone());
+                checking_data
             }
-        }) else {
-            debug!(%commitment, "No strong shards present to form checking data");
-            return Ok(None);
+            Entry::Occupied(entry) => entry.get().clone(),
         };
 
+        let cached_checked_shards = self.checked_shards.entry(commitment).or_default();
         let checked_shards = shards
             .into_iter()
             .filter_map(|s| {
@@ -238,18 +228,33 @@ where
                         // Any strong shards, at this point, were sent from the proposer.
                         // We use the reshard interface to produce our checked shard rather
                         // than taking two hops.
-                        S::reshard(&config, &commitment.inner(), index, shard)
-                            .map(|(_, checked, _)| checked)
-                            .ok()
+                        match cached_checked_shards.entry(index as usize) {
+                            Entry::Vacant(entry) => {
+                                let (_, checked, _) =
+                                    S::reshard(&config, &commitment.inner(), index, shard).ok()?;
+                                entry.insert(checked.clone());
+                                Some(checked)
+                            }
+                            Entry::Occupied(entry) => Some(entry.get().clone()),
+                        }
+                    }
+                    DistributionShard::Weak(re_shard) => {
+                        match cached_checked_shards.entry(index as usize) {
+                            Entry::Vacant(entry) => {
+                                let checked = S::check(
+                                    &config,
+                                    &commitment.inner(),
+                                    &checking_data,
+                                    index,
+                                    re_shard,
+                                )
+                                .ok()?;
+                                entry.insert(checked.clone());
+                                Some(checked)
+                            }
+                            Entry::Occupied(entry) => Some(entry.get().clone()),
+                        }
                     }
-                    DistributionShard::Weak(re_shard) => S::check(
-                        &config,
-                        &commitment.inner(),
-                        &checking_data,
-                        index,
-                        re_shard,
-                    )
-                    .ok(),
                 }
             })
             .collect::<Vec<_>>();
@@ -282,6 +287,8 @@ where
         );
 
         self.reconstructed_blocks.insert(commitment, block.clone());
+        self.checking_data.remove(&commitment);
+        self.checked_shards.remove(&commitment);
 
         // Notify any subscribers that have been waiting for this block to be reconstructed
         if let Some(mut sub) = self.block_subscriptions.remove(&commitment) {
@@ -363,6 +370,11 @@ where
     pub fn evict_block(&mut self, commitment: &CodingCommitment) {
         self.reconstructed_blocks.remove(commitment);
     }
+
+    /// Checks if a reconstructed block is present in the local cache.
+    pub fn has_block(&self, commitment: &CodingCommitment) -> bool {
+        self.reconstructed_blocks.contains_key(commitment)
+    }
 }
 
 #[cfg(test)]
diff --git a/consensus/src/marshal/ingress/mailbox.rs b/consensus/src/marshal/ingress/mailbox.rs
@@ -1,6 +1,6 @@
 use crate::{
     marshal::ingress::coding::types::CodedBlock,
-    threshold_simplex::types::{Activity, Finalization, Notarization},
+    threshold_simplex::types::{Activity, Finalization, Notarization, Notarize},
     types::Round,
     Block, Reporter,
 };
@@ -100,6 +100,10 @@ pub(crate) enum Message<V: Variant, B: Block, S: Scheme, P: PublicKey> {
     },
 
     // -------------------- Consensus Engine Messages --------------------
+    /// A notarize vote from the consensus engine.
+    Notarize {
+        notarization: Notarize<V, B::Commitment>,
+    },
     /// A notarization from the consensus engine.
     Notarization {
         /// The notarization.
@@ -254,6 +258,7 @@ impl<V: Variant, B: Block, S: Scheme, P: PublicKey> Reporter for Mailbox<V, B, S
 
     async fn report(&mut self, activity: Self::Activity) {
         let message = match activity {
+            Activity::Notarize(notarization) => Message::Notarize { notarization },
             Activity::Notarization(notarization) => Message::Notarization { notarization },
             Activity::Finalization(finalization) => Message::Finalization { finalization },
             _ => {

Original file line number	Diff line number	Diff line change
`@@ -461,7 +461,17 @@ where`
`461`	`461`	`}`
`462`	`462`	`}`
`463`	`463`	`}`
	`464`	`+ Message::Notarize { notarization } => {`
	`465`	`+ let commitment = notarization.proposal.payload;`
	`466`	`+ if !shards.has_block(&commitment) {`
	`467`	`+ let start = Instant::now();`
	`468`	`+ let _ = shards.try_reconstruct(commitment).await;`
	`469`	`+ let elapsed = start.elapsed();`
	`470`	`+ tracing::info!(?elapsed, "Attempted reconstruction");`
	`471`	`+ }`
	`472`	`+ }`
`464`	`473`	`Message::Notarization { notarization } => {`
	`474`	`+ tracing::warn!(?notarization, "Received notarization");`
`465`	`475`	`let round = notarization.round();`
`466`	`476`	`let commitment = notarization.proposal.payload;`
`467`	`477`