support subscriptions

Author: clabby

consensus/src/marshal/actor.rs

@@ -15,6 +15,7 @@ use crate::{
     Block, Reporter,
 };
 use commonware_codec::{Decode, Encode};
+use commonware_coding::reed_solomon::Chunk;
 use commonware_cryptography::{bls12381::primitives::variant::Variant, Hasher, PublicKey};
 use commonware_macros::select;
 use commonware_resolver::Resolver;
@@ -37,9 +38,17 @@ use tracing::{debug, info, warn};
 
 /// A struct that holds multiple subscriptions for a block.
 struct BlockSubscription<B: Block> {
-    // The subscribers that are waiting for the block
+    /// The subscribers that are waiting for the block
     subscribers: Vec<oneshot::Sender<B>>,
-    // Aborter that aborts the waiter future when dropped
+    /// Aborter that aborts the waiter future when dropped
+    _aborter: Aborter,
+}
+
+/// A struct that holds multiple subscriptions for a chunk.
+struct ChunkSubscription<H: Hasher> {
+    /// The subscribers that are waiting for the chunk
+    subscribers: Vec<oneshot::Sender<Chunk<H>>>,
+    /// Aborter that aborts the waiter future when dropped
     _aborter: Aborter,
 }
 
@@ -92,6 +101,8 @@ where
 
     // Outstanding subscriptions for blocks
     block_subscriptions: BTreeMap<B::Commitment, BlockSubscription<B>>,
+    // Outstanding subscriptions for chunks
+    chunk_subscriptions: BTreeMap<B::Commitment, ChunkSubscription<H>>,
 
     // ---------- Storage ----------
     // Prunable cache
@@ -232,6 +243,7 @@ where
                 codec_config: config.codec_config,
                 last_processed_round: Round::new(0, 0),
                 block_subscriptions: BTreeMap::new(),
+                chunk_subscriptions: BTreeMap::new(),
                 cache,
                 finalizations_by_height,
                 finalized_blocks,
@@ -282,7 +294,8 @@ where
             .spawn(|_| finalizer.run());
 
         // Create a local pool for waiter futures
-        let mut waiters = AbortablePool::<(B::Commitment, B)>::default();
+        let mut block_waiters = AbortablePool::<(B::Commitment, B)>::default();
+        let mut chunk_waiters = AbortablePool::<((B::Commitment, u16), Chunk<H>)>::default();
 
         // Handle messages
         loop {
@@ -295,7 +308,7 @@ where
             // Select messages
             select! {
                 // Handle waiter completions first
-                result = waiters.next_completed() => {
+                result = block_waiters.next_completed() => {
                     let Ok((commitment, block)) = result else {
                         continue; // Aborted future
                     };
@@ -364,6 +377,24 @@ where
                             let result = self.find_block(&mut shard_layer, commitment).await;
                             let _ = response.send(result);
                         }
+                        Message::SubscribeChunk { commitment, index, response } => {
+                            match self.chunk_subscriptions.entry(commitment) {
+                                Entry::Occupied(mut entry) => {
+                                    entry.get_mut().subscribers.push(response);
+                                }
+                                Entry::Vacant(entry) => {
+                                    let (tx, rx) = oneshot::channel();
+                                    shard_layer.subscribe_chunk(commitment, index, tx).await;
+                                    let aborter = chunk_waiters.push(async move {
+                                        ((commitment, index), rx.await.expect("shard subscriber closed"))
+                                    });
+                                    entry.insert(ChunkSubscription {
+                                        subscribers: vec![response],
+                                        _aborter: aborter,
+                                    });
+                                }
+                            }
+                        }
                         Message::Subscribe { round, commitment, response } => {
                             // Check for block locally
                             if let Some(block) = self.find_block(&mut shard_layer, commitment).await {
@@ -399,8 +430,8 @@ where
                                 }
                                 Entry::Vacant(entry) => {
                                     let (tx, rx) = oneshot::channel();
-                                    shard_layer.subscribe_prepared(commitment, tx).await.unwrap();
-                                    let aborter = waiters.push(async move {
+                                    shard_layer.subscribe_block(commitment, tx).await.expect("Reconstruction error not yet handled");
+                                    let aborter = block_waiters.push(async move {
                                         (commitment, rx.await.expect("shard subscriber closed"))
                                     });
                                     entry.insert(BlockSubscription {

consensus/src/marshal/ingress/coding.rs

@@ -22,6 +22,7 @@ use futures::channel::oneshot;
 use governor::clock::Clock as GClock;
 use rand::Rng;
 use std::{
+    collections::{btree_map::Entry, BTreeMap},
     fmt::Debug,
     num::{NonZero, NonZeroUsize},
     ops::Deref,
@@ -59,6 +60,16 @@ pub struct Config {
     pub buffer_pool: PoolRef,
 }
 
+/// A subscription for a block by its commitment.
+struct BlockSubscription<B: Block> {
+    subscribers: Vec<oneshot::Sender<B>>,
+}
+
+/// A subscription for a chunk by its commitment and index.
+struct ChunkSubscription<H: Hasher> {
+    subscribers: Vec<oneshot::Sender<Chunk<H>>>,
+}
+
 /// A layer that handles receiving erasure coded [Block]s from the [Actor](super::super::actor::Actor),
 /// broadcasting them to peers, and reassembling them from received [Shard]s.
 pub struct ShardLayer<E, P, B, H>
@@ -76,6 +87,12 @@ where
 
     /// Map of block digests -> coding commitments.
     digest_map: prunable::Archive<TwoCap, E, H::Digest, H::Digest>,
+
+    /// 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<E, P, B, H> ShardLayer<E, P, B, H>
@@ -110,6 +127,8 @@ where
             mailbox,
             block_codec_cfg,
             digest_map,
+            block_subscriptions: BTreeMap::new(),
+            chunk_subscriptions: BTreeMap::new(),
         }
     }
 
@@ -167,6 +186,16 @@ 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(subs) = self.chunk_subscriptions.get_mut(&(commitment, chunk.index)) {
+                for sub in subs.subscribers.drain(..) {
+                    let _ = sub.send(chunk.clone());
+                }
+                self.chunk_subscriptions.remove(&(commitment, chunk.index));
+            }
+        }
+
         if coded_chunks.len() < min as usize {
             // Not enough chunks to recover the block yet.
             debug!(
@@ -193,6 +222,14 @@ where
                 .await;
         }
 
+        // Attempt to resolve any open subscriptions for this block.
+        if let Some(subs) = self.block_subscriptions.get_mut(&commitment) {
+            for sub in subs.subscribers.drain(..) {
+                let _ = sub.send(block.clone());
+            }
+            self.block_subscriptions.remove(&commitment);
+        }
+
         info!(
             %commitment,
             digest = %block.digest(),
@@ -205,15 +242,59 @@ where
 
     /// Subscribes to a block by commitment with an externally prepared responder.
     ///
-    /// The responder will be sent the first message for a commitment when it is available; either
-    /// instantly (if cached) or when it is received from the network. The request can be canceled
-    /// by dropping the responder.
-    pub async fn subscribe_prepared(
+    /// The responder will be sent the block when it is available; either instantly (if cached)
+    /// or when it is received from the network. The request can be canceled by dropping the
+    /// responder.
+    pub async fn subscribe_block(
         &mut self,
-        _commitment: B::Commitment,
-        _responder: oneshot::Sender<B>,
+        commitment: B::Commitment,
+        responder: oneshot::Sender<B>,
     ) -> Result<(), ReconstructionError> {
-        todo!("Create subscription");
+        match self.block_subscriptions.entry(commitment) {
+            Entry::Vacant(entry) => {
+                entry.insert(BlockSubscription {
+                    subscribers: vec![responder],
+                });
+            }
+            Entry::Occupied(mut entry) => {
+                entry.get_mut().subscribers.push(responder);
+            }
+        }
+
+        // Try to reconstruct the block immediately in case we already have enough chunks.
+        self.try_reconstruct(commitment).await?;
+
+        Ok(())
+    }
+
+    /// 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)
+    /// or when it is received from the network. The request can be canceled by dropping the
+    /// responder.
+    pub async fn subscribe_chunk(
+        &mut self,
+        commitment: B::Commitment,
+        index: u16,
+        responder: oneshot::Sender<Chunk<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);
+            }
+        }
     }
 
     /// Puts a coding commitment in the store, keyed by digest and block height.

consensus/src/marshal/ingress/mailbox.rs

@@ -24,16 +24,27 @@ pub(crate) enum Message<V: Variant, B: Block, P: PublicKey, H: Hasher> {
         /// A channel to send the retrieved block.
         response: oneshot::Sender<Option<B>>,
     },
-    /// A request to retrieve a block by its digest.
+    /// A request to retrieve a block by its commitment.
+    ///
+    /// TODO: For a cleaner API, we should have this be retrieve-by-digest.
     Subscribe {
         /// The view in which the block was notarized. This is an optimization
         /// to help locate the block.
         round: Option<Round>,
-        /// The digest of the block to retrieve.
+        /// The coding commitment of the block to retrieve.
         commitment: B::Commitment,
         /// A channel to send the retrieved block.
         response: oneshot::Sender<B>,
     },
+    /// A request to retrieve an erasure coded chunk by its commitment and index.
+    SubscribeChunk {
+        /// The commitment of the chunk to retrieve.
+        commitment: B::Commitment,
+        /// The index of the chunk to retrieve.
+        index: u16,
+        /// A channel to send the retrieved chunk.
+        response: oneshot::Sender<Chunk<H>>,
+    },
     /// A request to broadcast an erasure coded block to all peers.
     Broadcast {
         /// The coding commitment of the block.
@@ -134,6 +145,36 @@ impl<V: Variant, B: Block, P: PublicKey, H: Hasher> Mailbox<V, B, P, H> {
         rx
     }
 
+    /// Subscribe chunk is a request to receive an erasure coded chunk by its commitment and index.
+    ///
+    /// If the chunk is found available locally, the chunk will be returned immediately.
+    ///
+    /// If the chunk is not available locally, the request will be registered and the caller will
+    /// be notified when the chunk is available. If the chunk is not part of a finalized block, it's
+    /// possible that it may never become available.
+    ///
+    /// The oneshot receiver should be dropped to cancel the subscription.
+    pub async fn subscribe_chunk(
+        &mut self,
+        commitment: B::Commitment,
+        index: u16,
+    ) -> oneshot::Receiver<Chunk<H>> {
+        let (tx, rx) = oneshot::channel();
+        if self
+            .sender
+            .send(Message::SubscribeChunk {
+                commitment,
+                index,
+                response: tx,
+            })
+            .await
+            .is_err()
+        {
+            error!("failed to send subscribe chunk message to actor: receiver dropped");
+        }
+        rx
+    }
+
     /// Broadcast indicates that a block should be sent to all peers.
     pub async fn broadcast(
         &mut self,

consensus/src/marshal/mod.rs

@@ -377,7 +377,7 @@ mod tests {
         }
     }
 
-    #[test_traced("WARN")]
+    #[test_traced("DEBUG")]
     fn test_finalize_bad_links() {
         for seed in 0..5 {
             let result1 = finalize(seed, UNRELIABLE_LINK);
@@ -468,6 +468,10 @@ mod tests {
                         .await;
                 }
 
+                // Wait for the block chunks to be delivered; Before making a notarization,
+                // the chunks must be present.
+                context.sleep(link.latency + link.jitter).await;
+
                 // Notarize block by the validator that broadcasted it
                 let notarization = make_notarization(proposal.clone(), &shares, QUORUM);
                 actor
@@ -623,9 +627,6 @@ mod tests {
             let (commitment1, config1, chunks1) = shard(&block1, &peers);
             let (commitment2, config2, chunks2) = shard(&block2, &peers);
 
-            actor.broadcast(commitment1, config1, chunks1).await;
-            actor.broadcast(commitment2, config2, chunks2).await;
-
             let sub1_rx = actor
                 .subscribe(Some(Round::from((0, 1))), commitment1)
                 .await;
@@ -636,6 +637,9 @@ mod tests {
                 .subscribe(Some(Round::from((0, 1))), commitment1)
                 .await;
 
+            actor.broadcast(commitment1, config1, chunks1).await;
+            actor.broadcast(commitment2, config2, chunks2).await;
+
             for (view, block) in [(1, block1.clone()), (2, block2.clone())] {
                 let proposal = Proposal {
                     round: Round::new(0, view),