fix shard sub

clabby · clabby · commit ad6e255d06b9 · 2025-09-24T14:27:40.000-07:00
diff --git a/consensus/src/marshal/actor.rs b/consensus/src/marshal/actor.rs
@@ -9,7 +9,7 @@ use super::{
     },
 };
 use crate::{
-    marshal::ingress::coding::{CodedBlock, ShardLayer},
+    marshal::ingress::coding::{CodedBlock, Shard, ShardLayer},
     threshold_simplex::types::{Finalization, Notarization},
     types::Round,
     Block, Reporter,
@@ -47,9 +47,13 @@ struct BlockSubscription<B: Block> {
 }
 
 /// A struct that holds multiple subscriptions for a chunk.
-struct ChunkSubscription<H: Hasher> {
+struct ChunkSubscription<B, H>
+where
+    B: Block<Digest = H::Digest, Commitment = H::Digest>,
+    H: Hasher,
+{
     /// The subscribers that are waiting for the chunk
-    subscribers: Vec<oneshot::Sender<Chunk<H>>>,
+    subscribers: Vec<oneshot::Sender<Shard<CodedBlock<B, H>, H>>>,
     /// Aborter that aborts the waiter future when dropped
     _aborter: Aborter,
 }
@@ -104,7 +108,7 @@ where
     // Outstanding subscriptions for blocks
     block_subscriptions: BTreeMap<B::Commitment, BlockSubscription<B>>,
     // Outstanding subscriptions for chunks
-    chunk_subscriptions: BTreeMap<B::Commitment, ChunkSubscription<H>>,
+    chunk_subscriptions: BTreeMap<B::Commitment, ChunkSubscription<B, H>>,
 
     // ---------- Storage ----------
     // Prunable cache
@@ -297,7 +301,8 @@ where
 
         // Create a local pool for waiter futures
         let mut block_waiters = AbortablePool::<(B::Commitment, CodedBlock<B, H>)>::default();
-        let mut chunk_waiters = AbortablePool::<((B::Commitment, u16), Chunk<H>)>::default();
+        let mut chunk_waiters =
+            AbortablePool::<((B::Commitment, u16), Shard<CodedBlock<B, H>, H>)>::default();
 
         // Handle messages
         loop {
@@ -381,6 +386,12 @@ where
                             let _ = response.send(result.map(CodedBlock::take_inner));
                         }
                         Message::SubscribeChunk { commitment, index, response } => {
+                            // Check for chunk locally
+                            if let Some(shard) = shard_layer.get_chunk(commitment, index).await {
+                                let _ = response.send(shard);
+                                continue;
+                            }
+
                             match self.chunk_subscriptions.entry(commitment) {
                                 Entry::Occupied(mut entry) => {
                                     entry.get_mut().subscribers.push(response);
diff --git a/consensus/src/marshal/ingress/coding.rs b/consensus/src/marshal/ingress/coding.rs
@@ -9,7 +9,7 @@
 
 use crate::Block;
 use commonware_broadcast::{buffered, Broadcaster};
-use commonware_codec::{Encode, EncodeSize, Error as CodecError, Read, ReadExt, Write};
+use commonware_codec::{EncodeSize, Error as CodecError, FixedSize, Read, ReadExt, Write};
 use commonware_coding::reed_solomon::{self, decode, Chunk, Error as ReedSolomonError};
 use commonware_cryptography::{Committable, Digestible, Hasher, PublicKey};
 use commonware_p2p::Recipients;
@@ -60,9 +60,6 @@ where
 
     /// Open subscriptions for blocks by commitment.
     block_subscriptions: BTreeMap<B::Commitment, BlockSubscription<B>>,
-
-    /// Open subscriptions for chunks by commitment and index.
-    chunk_subscriptions: BTreeMap<(B::Commitment, u16), ChunkSubscription<H>>,
 }
 
 impl<P, B, H> ShardLayer<P, B, H>
@@ -77,7 +74,6 @@ where
             mailbox,
             block_codec_cfg,
             block_subscriptions: BTreeMap::new(),
-            chunk_subscriptions: BTreeMap::new(),
         }
     }
 
@@ -135,15 +131,6 @@ where
             .map(|c| c.chunk)
             .collect::<Vec<_>>();
 
-        // Attempt to resolve any open subscriptions for the chunks we have.
-        for chunk in coded_chunks.iter() {
-            if let Some(mut subs) = self.chunk_subscriptions.remove(&(commitment, chunk.index)) {
-                for sub in subs.subscribers.drain(..) {
-                    let _ = sub.send(chunk.clone());
-                }
-            }
-        }
-
         if coded_chunks.len() < min as usize {
             // Not enough chunks to recover the block yet.
             debug!(
@@ -206,6 +193,23 @@ where
         Ok(())
     }
 
+    pub async fn get_chunk(
+        &mut self,
+        commitment: B::Commitment,
+        index: u16,
+    ) -> Option<Shard<B, H>> {
+        let mut buf = vec![0u8; <H::Digest as FixedSize>::SIZE + 2];
+        buf[..<H::Digest as FixedSize>::SIZE].copy_from_slice(commitment.as_ref());
+        buf[<H::Digest as FixedSize>::SIZE..].copy_from_slice(index.to_le_bytes().as_ref());
+        let index_hash = H::hash(buf.as_ref());
+        self.mailbox
+            .get(None, commitment, Some(index_hash))
+            .await
+            .iter()
+            .cloned()
+            .next()
+    }
+
     /// Subscribes to a chunk by commitment and index with an externally prepared responder.
     ///
     /// The responder will be sent the chunk when it is available; either instantly (if cached)
@@ -215,25 +219,15 @@ where
         &mut self,
         commitment: B::Commitment,
         index: u16,
-        responder: oneshot::Sender<Chunk<H>>,
+        responder: oneshot::Sender<Shard<B, H>>,
     ) {
-        let available_chunks = self.mailbox.get(None, commitment, None).await;
-
-        if let Some(shard) = available_chunks.iter().find(|s| s.chunk.index == index) {
-            let _ = responder.send(shard.chunk.clone());
-            return;
-        }
-
-        match self.chunk_subscriptions.entry((commitment, index)) {
-            Entry::Vacant(entry) => {
-                entry.insert(ChunkSubscription {
-                    subscribers: vec![responder],
-                });
-            }
-            Entry::Occupied(mut entry) => {
-                entry.get_mut().subscribers.push(responder);
-            }
-        }
+        let mut buf = vec![0u8; <H::Digest as FixedSize>::SIZE + 2];
+        buf[..<H::Digest as FixedSize>::SIZE].copy_from_slice(commitment.as_ref());
+        buf[<H::Digest as FixedSize>::SIZE..].copy_from_slice(index.to_le_bytes().as_ref());
+        let index_hash = H::hash(buf.as_ref());
+        self.mailbox
+            .subscribe_prepared(None, commitment, Some(index_hash), responder)
+            .await;
     }
 }
 
@@ -326,9 +320,11 @@ where
     type Digest = H::Digest;
 
     fn digest(&self) -> Self::Digest {
-        // NOTE: This is a lil weird; only doing this to namespace the shard within the buffered mailbox, such that
-        // shards from separate validators can be enqueued without replacing each other.
-        H::hash(self.chunk.encode().as_ref())
+        let mut buf = vec![0u8; <H::Digest as FixedSize>::SIZE + 2];
+        buf[..<H::Digest as FixedSize>::SIZE].copy_from_slice(self.commitment.as_ref());
+        buf[<H::Digest as FixedSize>::SIZE..]
+            .copy_from_slice(self.chunk.index.to_le_bytes().as_ref());
+        H::hash(buf.as_ref())
     }
 }
 
diff --git a/consensus/src/marshal/ingress/mailbox.rs b/consensus/src/marshal/ingress/mailbox.rs
@@ -1,4 +1,5 @@
 use crate::{
+    marshal::ingress::coding::{CodedBlock, Shard},
     threshold_simplex::types::{Activity, Finalization, Finalize, Notarization, Notarize},
     types::Round,
     Block, Reporter,
@@ -15,7 +16,12 @@ use tracing::error;
 ///
 /// These messages are sent from the consensus engine and other parts of the
 /// system to drive the state of the marshal.
-pub(crate) enum Message<V: Variant, B: Block, P: PublicKey, H: Hasher> {
+pub(crate) enum Message<
+    V: Variant,
+    B: Block<Digest = H::Digest, Commitment = H::Digest>,
+    P: PublicKey,
+    H: Hasher,
+> {
     // -------------------- Application Messages --------------------
     /// A request to retrieve a block by its commitment.
     Get {
@@ -41,7 +47,7 @@ pub(crate) enum Message<V: Variant, B: Block, P: PublicKey, H: Hasher> {
         /// The index of the chunk to retrieve.
         index: u16,
         /// A channel to send the retrieved chunk.
-        response: oneshot::Sender<Chunk<H>>,
+        response: oneshot::Sender<Shard<CodedBlock<B, H>, H>>,
     },
     /// A request to broadcast an erasure coded block to all peers.
     Broadcast {
@@ -78,11 +84,18 @@ pub(crate) enum Message<V: Variant, B: Block, P: PublicKey, H: Hasher> {
 
 /// A mailbox for sending messages to the marshal [Actor](super::super::actor::Actor).
 #[derive(Clone)]
-pub struct Mailbox<V: Variant, B: Block, P: PublicKey, H: Hasher> {
+pub struct Mailbox<
+    V: Variant,
+    B: Block<Digest = H::Digest, Commitment = H::Digest>,
+    P: PublicKey,
+    H: Hasher,
+> {
     sender: mpsc::Sender<Message<V, B, P, H>>,
 }
 
-impl<V: Variant, B: Block, P: PublicKey, H: Hasher> Mailbox<V, B, P, H> {
+impl<V: Variant, B: Block<Digest = H::Digest, Commitment = H::Digest>, P: PublicKey, H: Hasher>
+    Mailbox<V, B, P, H>
+{
     /// Creates a new mailbox.
     pub(crate) fn new(sender: mpsc::Sender<Message<V, B, P, H>>) -> Self {
         Self { sender }
@@ -149,7 +162,7 @@ impl<V: Variant, B: Block, P: PublicKey, H: Hasher> Mailbox<V, B, P, H> {
         &mut self,
         commitment: B::Commitment,
         index: u16,
-    ) -> oneshot::Receiver<Chunk<H>> {
+    ) -> oneshot::Receiver<Shard<CodedBlock<B, H>, H>> {
         let (tx, rx) = oneshot::channel();
         if self
             .sender
@@ -188,7 +201,9 @@ impl<V: Variant, B: Block, P: PublicKey, H: Hasher> Mailbox<V, B, P, H> {
     }
 }
 
-impl<V: Variant, B: Block, P: PublicKey, H: Hasher> Reporter for Mailbox<V, B, P, H> {
+impl<V: Variant, B: Block<Digest = H::Digest, Commitment = H::Digest>, P: PublicKey, H: Hasher>
+    Reporter for Mailbox<V, B, P, H>
+{
     type Activity = Activity<V, B::Commitment>;
 
     async fn report(&mut self, activity: Self::Activity) {
