[consensus/marshal] Cheaper CodedBlock::clone (#3708)

Author: clabby

consensus/src/marshal/coding/shards/engine.rs

@@ -174,7 +174,6 @@ use rand::Rng;
 use std::{
     collections::{BTreeMap, VecDeque},
     num::NonZeroUsize,
-    sync::Arc,
 };
 use thiserror::Error;
 use tracing::{debug, warn};
@@ -331,8 +330,7 @@ where
     /// An ephemeral cache of reconstructed blocks, keyed by commitment.
     ///
     /// These blocks are evicted after a durability signal from the marshal.
-    /// Wrapped in [`Arc`] to enable cheap cloning when serving multiple subscribers.
-    reconstructed_blocks: BTreeMap<Commitment, Arc<CodedBlock<B, C, H>>>,
+    reconstructed_blocks: BTreeMap<Commitment, CodedBlock<B, C, H>>,
 
     /// Open subscriptions for assigned shard verification for the keyed
     /// [`Commitment`].
@@ -350,7 +348,7 @@ where
     /// the keyed [`Commitment`].
     #[allow(clippy::type_complexity)]
     block_subscriptions:
-        BTreeMap<BlockSubscriptionKey<B::Digest>, Vec<oneshot::Sender<Arc<CodedBlock<B, C, H>>>>>,
+        BTreeMap<BlockSubscriptionKey<B::Digest>, Vec<oneshot::Sender<CodedBlock<B, C, H>>>>,
 
     /// Metrics for the shard engine.
     metrics: ShardMetrics<P>,
@@ -580,9 +578,9 @@ where
     fn try_reconstruct(
         &mut self,
         commitment: Commitment,
-    ) -> Result<Option<Arc<CodedBlock<B, C, H>>>, Error<C>> {
+    ) -> Result<Option<CodedBlock<B, C, H>>, Error<C>> {
         if let Some(block) = self.reconstructed_blocks.get(&commitment) {
-            return Ok(Some(Arc::clone(block)));
+            return Ok(Some(block.clone()));
         }
         let Some(state) = self.state.get_mut(&commitment) else {
             return Ok(None);
@@ -635,8 +633,8 @@ where
 
         // Construct a coding block with a _trusted_ commitment. `S::decode` verified the blob's
         // integrity against the commitment, so shards can be lazily re-constructed if need be.
-        let block = Arc::new(CodedBlock::new_trusted(inner, commitment));
-        self.cache_block(Arc::clone(&block));
+        let block = CodedBlock::new_trusted(inner, commitment);
+        self.cache_block(block.clone());
         self.metrics.blocks_reconstructed_total.inc();
         Ok(Some(block))
     }
@@ -744,10 +742,9 @@ where
     }
 
     /// Cache a block and notify all subscribers waiting on it.
-    fn cache_block(&mut self, block: Arc<CodedBlock<B, C, H>>) {
+    fn cache_block(&mut self, block: CodedBlock<B, C, H>) {
         let commitment = block.commitment();
-        self.reconstructed_blocks
-            .insert(commitment, Arc::clone(&block));
+        self.reconstructed_blocks.insert(commitment, block.clone());
         self.notify_block_subscribers(block);
     }
 
@@ -825,7 +822,6 @@ where
         }
 
         // Cache the block so we don't have to reconstruct it again.
-        let block = Arc::new(block);
         self.cache_block(block);
 
         // Local proposals bypass reconstruction, so shard subscribers waiting
@@ -937,7 +933,7 @@ where
     fn handle_block_subscription(
         &mut self,
         key: BlockSubscriptionKey<B::Digest>,
-        response: oneshot::Sender<Arc<CodedBlock<B, C, H>>>,
+        response: oneshot::Sender<CodedBlock<B, C, H>>,
     ) {
         let block = match key {
             BlockSubscriptionKey::Commitment(commitment) => {
@@ -951,7 +947,7 @@ where
 
         // Answer immediately if we have the block cached.
         if let Some(block) = block {
-            response.send_lossy(Arc::clone(block));
+            response.send_lossy(block.clone());
             return;
         }
 
@@ -975,7 +971,7 @@ where
     }
 
     /// Notifies and cleans up any subscriptions for a reconstructed block.
-    fn notify_block_subscribers(&mut self, block: Arc<CodedBlock<B, C, H>>) {
+    fn notify_block_subscribers(&mut self, block: CodedBlock<B, C, H>) {
         let commitment = block.commitment();
         let digest = block.digest();
 
@@ -985,7 +981,7 @@ where
             .remove(&BlockSubscriptionKey::Commitment(commitment))
         {
             for subscriber in subscribers.drain(..) {
-                subscriber.send_lossy(Arc::clone(&block));
+                subscriber.send_lossy(block.clone());
             }
         }
 
@@ -995,7 +991,7 @@ where
             .remove(&BlockSubscriptionKey::Digest(digest))
         {
             for subscriber in subscribers.drain(..) {
-                subscriber.send_lossy(Arc::clone(&block));
+                subscriber.send_lossy(block.clone());
             }
         }
     }
@@ -1541,7 +1537,10 @@ mod tests {
         future::Future,
         marker::PhantomData,
         num::NonZeroU32,
-        sync::atomic::{AtomicIsize, Ordering},
+        sync::{
+            atomic::{AtomicIsize, Ordering},
+            Arc,
+        },
         time::Duration,
     };
 

consensus/src/marshal/coding/shards/mailbox.rs

@@ -8,7 +8,6 @@ use crate::{
 use commonware_coding::Scheme as CodingScheme;
 use commonware_cryptography::{Hasher, PublicKey};
 use commonware_utils::channel::{fallible::AsyncFallibleExt, mpsc, oneshot};
-use std::sync::Arc;
 
 /// A message that can be sent to the coding [`Engine`].
 ///
@@ -41,14 +40,14 @@ where
         /// The [`Commitment`] of the block to get.
         commitment: Commitment,
         /// The response channel.
-        response: oneshot::Sender<Option<Arc<CodedBlock<B, C, H>>>>,
+        response: oneshot::Sender<Option<CodedBlock<B, C, H>>>,
     },
     /// A request to get a reconstructed block by its digest, if available.
     GetByDigest {
         /// The digest of the block to get.
         digest: B::Digest,
         /// The response channel.
-        response: oneshot::Sender<Option<Arc<CodedBlock<B, C, H>>>>,
+        response: oneshot::Sender<Option<CodedBlock<B, C, H>>>,
     },
     /// A request to open a subscription for assigned shard verification.
     ///
@@ -72,15 +71,15 @@ where
         /// The block's digest.
         commitment: Commitment,
         /// The response channel.
-        response: oneshot::Sender<Arc<CodedBlock<B, C, H>>>,
+        response: oneshot::Sender<CodedBlock<B, C, H>>,
     },
     /// A request to open a subscription for the reconstruction of a [`CodedBlock`]
     /// by its digest.
     SubscribeByDigest {
         /// The block's digest.
         digest: B::Digest,
         /// The response channel.
-        response: oneshot::Sender<Arc<CodedBlock<B, C, H>>>,
+        response: oneshot::Sender<CodedBlock<B, C, H>>,
     },
     /// A request to prune all caches at and below the given commitment.
     Prune {
@@ -132,7 +131,7 @@ where
     }
 
     /// Request a reconstructed block by its [`Commitment`].
-    pub async fn get(&self, commitment: Commitment) -> Option<Arc<CodedBlock<B, C, H>>> {
+    pub async fn get(&self, commitment: Commitment) -> Option<CodedBlock<B, C, H>> {
         self.sender
             .request(|tx| Message::GetByCommitment {
                 commitment,
@@ -143,7 +142,7 @@ where
     }
 
     /// Request a reconstructed block by its digest.
-    pub async fn get_by_digest(&self, digest: B::Digest) -> Option<Arc<CodedBlock<B, C, H>>> {
+    pub async fn get_by_digest(&self, digest: B::Digest) -> Option<CodedBlock<B, C, H>> {
         self.sender
             .request(|tx| Message::GetByDigest {
                 digest,
@@ -180,7 +179,7 @@ where
     pub async fn subscribe(
         &self,
         commitment: Commitment,
-    ) -> oneshot::Receiver<Arc<CodedBlock<B, C, H>>> {
+    ) -> oneshot::Receiver<CodedBlock<B, C, H>> {
         let (responder, receiver) = oneshot::channel();
         let msg = Message::SubscribeByCommitment {
             commitment,
@@ -194,7 +193,7 @@ where
     pub async fn subscribe_by_digest(
         &self,
         digest: B::Digest,
-    ) -> oneshot::Receiver<Arc<CodedBlock<B, C, H>>> {
+    ) -> oneshot::Receiver<CodedBlock<B, C, H>> {
         let (responder, receiver) = oneshot::channel();
         let msg = Message::SubscribeByDigest {
             digest,

consensus/src/marshal/coding/types.rs

@@ -9,7 +9,7 @@ use commonware_coding::{Config as CodingConfig, Scheme};
 use commonware_cryptography::{Committable, Digestible, Hasher};
 use commonware_parallel::{Sequential, Strategy};
 use commonware_utils::{Faults, N3f1, NZU16};
-use std::marker::PhantomData;
+use std::{marker::PhantomData, sync::Arc};
 
 /// A broadcastable shard of erasure coded data, including the coding commitment and
 /// the configuration used to code the data.
@@ -142,7 +142,7 @@ where
 #[derive(Debug)]
 pub struct CodedBlock<B: Block, C: Scheme, H: Hasher> {
     /// The inner block type.
-    inner: B,
+    inner: Arc<B>,
     /// The erasure coding configuration.
     config: CodingConfig,
     /// The erasure coding commitment.
@@ -152,7 +152,7 @@ pub struct CodedBlock<B: Block, C: Scheme, H: Hasher> {
     /// These shards are optional to enable lazy construction. If the block is
     /// constructed with [`Self::new_trusted`], the shards are computed lazily
     /// via [`Self::shards`].
-    shards: Option<Vec<C::Shard>>,
+    shards: Option<Arc<[C::Shard]>>,
     /// Phantom data for the hasher.
     _hasher: PhantomData<H>,
 }
@@ -175,18 +175,18 @@ impl<B: Block, C: Scheme, H: Hasher> CodedBlock<B, C, H> {
     pub fn new(inner: B, config: CodingConfig, strategy: &impl Strategy) -> Self {
         let (commitment, shards) = Self::encode(&inner, config, strategy);
         Self {
-            inner,
+            inner: Arc::new(inner),
             config,
             commitment,
-            shards: Some(shards),
+            shards: Some(shards.into()),
             _hasher: PhantomData,
         }
     }
 
     /// Create a new [`CodedBlock`] from a [`Block`] and trusted [`Commitment`].
     pub fn new_trusted(inner: B, commitment: Commitment) -> Self {
         Self {
-            inner,
+            inner: Arc::new(inner),
             config: commitment.config(),
             commitment: commitment.root(),
             shards: None,
@@ -213,7 +213,7 @@ impl<B: Block, C: Scheme, H: Hasher> CodedBlock<B, C, H> {
                     "coded block constructed with trusted commitment does not match commitment"
                 );
 
-                self.shards = Some(shards);
+                self.shards = Some(shards.into());
                 self.shards.as_ref().unwrap()
             }
         }
@@ -232,28 +232,28 @@ impl<B: Block, C: Scheme, H: Hasher> CodedBlock<B, C, H> {
     }
 
     /// Returns a reference to the inner [`Block`].
-    pub const fn inner(&self) -> &B {
+    pub fn inner(&self) -> &B {
         &self.inner
     }
 
     /// Takes the inner [`Block`].
     pub fn into_inner(self) -> B {
-        self.inner
+        Arc::unwrap_or_clone(self.inner)
     }
 }
 
-impl<B: CertifiableBlock, C: Scheme, H: Hasher> From<CodedBlock<B, C, H>>
+impl<B: CertifiableBlock + Clone, C: Scheme, H: Hasher> From<CodedBlock<B, C, H>>
     for StoredCodedBlock<B, C, H>
 {
     fn from(block: CodedBlock<B, C, H>) -> Self {
         Self::new(block)
     }
 }
 
-impl<B: Block + Clone, C: Scheme, H: Hasher> Clone for CodedBlock<B, C, H> {
+impl<B: Block, C: Scheme, H: Hasher> Clone for CodedBlock<B, C, H> {
     fn clone(&self) -> Self {
         Self {
-            inner: self.inner.clone(),
+            inner: Arc::clone(&self.inner),
             config: self.config,
             commitment: self.commitment,
             shards: self.shards.clone(),
@@ -315,10 +315,10 @@ impl<B: Block, C: Scheme, H: Hasher> Read for CodedBlock<B, C, H> {
             })?;
 
         Ok(Self {
-            inner,
+            inner: Arc::new(inner),
             config,
             commitment,
-            shards: Some(shards),
+            shards: Some(shards.into()),
             _hasher: PhantomData,
         })
     }
@@ -380,15 +380,15 @@ pub struct StoredCodedBlock<B: Block, C: Scheme, H: Hasher> {
     _scheme: PhantomData<(C, H)>,
 }
 
-impl<B: CertifiableBlock, C: Scheme, H: Hasher> StoredCodedBlock<B, C, H> {
+impl<B: CertifiableBlock + Clone, C: Scheme, H: Hasher> StoredCodedBlock<B, C, H> {
     /// Create a [`StoredCodedBlock`] from a verified [`CodedBlock`].
     ///
     /// The caller must ensure the [`CodedBlock`] has been properly verified
     /// (i.e., its commitment was computed or validated against a trusted source).
     pub fn new(block: CodedBlock<B, C, H>) -> Self {
         Self {
             commitment: block.commitment(),
-            inner: block.inner,
+            inner: block.into_inner(),
             _scheme: PhantomData,
         }
     }
@@ -621,6 +621,24 @@ mod test {
         assert!(coded_block == decoded);
     }
 
+    #[test]
+    fn test_coded_block_clone_shares_storage() {
+        const CONFIG: CodingConfig = CodingConfig {
+            minimum_shards: NZU16!(1),
+            extra_shards: NZU16!(2),
+        };
+
+        let block = Block::new::<Sha256>((), Sha256::hash(b"parent"), Height::new(42), 1_234_567);
+        let coded_block = CodedBlock::<Block, RS, H>::new(block, CONFIG, &Sequential);
+        let cloned = coded_block.clone();
+
+        assert!(Arc::ptr_eq(&coded_block.inner, &cloned.inner));
+        assert!(Arc::ptr_eq(
+            coded_block.shards.as_ref().unwrap(),
+            cloned.shards.as_ref().unwrap()
+        ));
+    }
+
     #[test]
     fn test_stored_coded_block_codec_roundtrip() {
         const CONFIG: CodingConfig = CodingConfig {

consensus/src/marshal/coding/variant.rs

@@ -14,7 +14,6 @@ use commonware_coding::Scheme as CodingScheme;
 use commonware_cryptography::{Committable, Digestible, Hasher, PublicKey};
 use commonware_p2p::Recipients;
 use commonware_utils::channel::oneshot;
-use std::sync::Arc;
 
 /// The coding variant of Marshal, which uses erasure coding for block dissemination.
 ///
@@ -68,30 +67,29 @@ where
     P: PublicKey,
 {
     type PublicKey = P;
-    type CachedBlock = Arc<CodedBlock<B, C, H>>;
 
     async fn find_by_digest(
         &self,
         digest: <CodedBlock<B, C, H> as Digestible>::Digest,
-    ) -> Option<Self::CachedBlock> {
+    ) -> Option<CodedBlock<B, C, H>> {
         self.get_by_digest(digest).await
     }
 
-    async fn find_by_commitment(&self, commitment: Commitment) -> Option<Self::CachedBlock> {
+    async fn find_by_commitment(&self, commitment: Commitment) -> Option<CodedBlock<B, C, H>> {
         self.get(commitment).await
     }
 
     async fn subscribe_by_digest(
         &self,
         digest: <CodedBlock<B, C, H> as Digestible>::Digest,
-    ) -> oneshot::Receiver<Self::CachedBlock> {
+    ) -> oneshot::Receiver<CodedBlock<B, C, H>> {
         self.subscribe_by_digest(digest).await
     }
 
     async fn subscribe_by_commitment(
         &self,
         commitment: Commitment,
-    ) -> oneshot::Receiver<Self::CachedBlock> {
+    ) -> oneshot::Receiver<CodedBlock<B, C, H>> {
         self.subscribe(commitment).await
     }