use prunable::Archive for mappings

clabby · clabby · commit dbaa491f1fe7 · 2025-09-22T20:08:43.000-07:00
diff --git a/consensus/src/marshal/actor.rs b/consensus/src/marshal/actor.rs
@@ -246,7 +246,7 @@ impl<
     pub fn start<R>(
         mut self,
         application: impl Reporter<Activity = B>,
-        shards: ShardLayer<P, B, H>,
+        shards: ShardLayer<E, P, B, H>,
         resolver: (mpsc::Receiver<handler::Message<B>>, R),
     ) -> Handle<()>
     where
@@ -259,7 +259,7 @@ impl<
     async fn run<R>(
         mut self,
         application: impl Reporter<Activity = B>,
-        mut shard_layer: ShardLayer<P, B, H>,
+        mut shard_layer: ShardLayer<E, P, B, H>,
         (mut resolver_rx, mut resolver): (mpsc::Receiver<handler::Message<B>>, R),
     ) where
         R: Resolver<Key = handler::Request<B>>,
@@ -327,13 +327,13 @@ impl<
 
                             // Block on waiting for the block to be reconstructed; We cannot move forward without
                             // having the block digest.
-                            while !shard_layer.has_digest(&commitment) {
+                            while !shard_layer.has_digest(&commitment).await {
                                 dbg!("spinning");
                                 shard_layer.try_reconstruct(commitment).await.expect("Reconstruction error not yet handled");
                             }
 
                             // Request the block digest for the block corresponding to the coding commitment.
-                            let digest = shard_layer.get_digest(&commitment).unwrap();
+                            let digest = shard_layer.get_digest(&commitment).await.unwrap();
 
                             // Store notarization by view
                             self.cache.put_notarization(round, digest, notarization.clone()).await;
@@ -362,13 +362,13 @@ impl<
 
                             // Block on waiting for the block to be reconstructed; We cannot move forward without
                             // having the block digest.
-                            while !shard_layer.has_digest(&commitment) {
+                            while !shard_layer.has_digest(&commitment).await {
                                 dbg!("spinning");
                                 shard_layer.try_reconstruct(commitment).await.expect("Reconstruction error not yet handled");
                             }
 
                             // Request the block digest for the block corresponding to the coding commitment.
-                            let digest = shard_layer.get_digest(&commitment).unwrap();
+                            let digest = shard_layer.get_digest(&commitment).await.unwrap();
 
                             self.cache.put_finalization(round, digest, finalization.clone()).await;
 
@@ -465,6 +465,7 @@ impl<
 
                                 // Prune archives
                                 self.cache.prune(prune_round).await;
+                                shard_layer.prune(height).await;
 
                                 // Update the last processed round
                                 let round = finalization.round();
@@ -558,7 +559,7 @@ impl<
 
                                     // Get block
                                     let commitment = notarization.proposal.payload;
-                                    let digest = match shard_layer.get_digest(&commitment) {
+                                    let digest = match shard_layer.get_digest(&commitment).await {
                                         Some(digest) => digest,
                                         None => {
                                             debug!(?commitment, "notarized block missing commitment mapping on request");
@@ -754,11 +755,15 @@ impl<
     /// Looks for a block anywhere in local storage.
     async fn find_block(
         &mut self,
-        shards: &mut ShardLayer<P, B, H>,
+        shards: &mut ShardLayer<E, P, B, H>,
         commitment: B::Commitment,
     ) -> Option<B> {
         // Check shard layer.
-        if let Some(block) = shards.get(commitment) {
+        if let Some(block) = shards
+            .block_by_digest(commitment)
+            .await
+            .expect("reconstruction error not yet handled")
+        {
             return Some(block);
         }
         // Check verified / notarized blocks via cache manager.
diff --git a/consensus/src/marshal/ingress/coding.rs b/consensus/src/marshal/ingress/coding.rs
@@ -13,8 +13,19 @@ use commonware_codec::{Encode, EncodeSize, Error as CodecError, Read, ReadExt, W
 use commonware_coding::reed_solomon::{decode, Chunk, Error as ReedSolomonError};
 use commonware_cryptography::{Committable, Digestible, Hasher, PublicKey};
 use commonware_p2p::Recipients;
+use commonware_runtime::{buffer::PoolRef, Clock, Metrics, Spawner, Storage};
+use commonware_storage::{
+    archive::{prunable, Archive, Identifier},
+    translator::TwoCap,
+};
 use futures::channel::oneshot;
-use std::{collections::HashMap, fmt::Debug, ops::Deref};
+use governor::clock::Clock as GClock;
+use rand::Rng;
+use std::{
+    fmt::Debug,
+    num::{NonZero, NonZeroUsize},
+    ops::Deref,
+};
 use thiserror::Error;
 use tracing::{debug, info};
 
@@ -30,11 +41,29 @@ pub enum ReconstructionError {
     Codec(#[from] CodecError),
 }
 
+/// Storage configuration for the [ShardLayer].
+pub struct Config {
+    /// Namespace prefix for the underlying storage partitions.
+    pub partition_prefix: String,
+
+    /// Number of items per section in the [prunable::Archive].
+    pub items_per_section: NonZero<u64>,
+
+    /// Rpelay buffer size for the [prunable::Archive].
+    pub replay_buffer: NonZeroUsize,
+
+    /// Write buffer size for the [prunable::Archive].
+    pub write_buffer: NonZeroUsize,
+
+    /// Buffer pool for the [prunable::Archive].
+    pub buffer_pool: PoolRef,
+}
+
 /// 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.
-#[derive(Clone)]
-pub struct ShardLayer<P, B, H>
+pub struct ShardLayer<E, P, B, H>
 where
+    E: Rng + Spawner + Metrics + Clock + GClock + Storage,
     P: PublicKey,
     B: Block<Digest = H::Digest, Commitment = H::Digest>,
     H: Hasher,
@@ -45,30 +74,50 @@ where
     /// [`Read`] configuration for the block type.
     block_codec_cfg: B::Cfg,
 
-    /// A map of coding commitments to block digests.
-    ///
-    /// TODO: This map has no durability nor pruning; just for testing / getting things working...
-    commitment_map: HashMap<B::Commitment, B::Digest>,
+    /// Map of block digests -> coding commitments.
+    commitment_map: prunable::Archive<TwoCap, E, H::Digest, H::Digest>,
 
-    /// A map of block digest to reconstructed blocks.
-    ///
-    /// TODO: This map has no durability nor pruning; just for testing / getting things working...
-    reconstruction_cache: HashMap<B::Digest, B>,
+    /// Map of coding commitments -> block digests.
+    digest_map: prunable::Archive<TwoCap, E, H::Digest, H::Digest>,
 }
 
-impl<P, B, H> ShardLayer<P, B, H>
+impl<E, P, B, H> ShardLayer<E, P, B, H>
 where
+    E: Rng + Spawner + Metrics + Clock + GClock + Storage,
     P: PublicKey,
     B: Block<Digest = H::Digest, Commitment = H::Digest> + Debug,
     H: Hasher,
 {
     /// Create a new [ShardLayer] with the given buffered mailbox.
-    pub fn new(mailbox: buffered::Mailbox<P, Shard<B, H>>, cfg: B::Cfg) -> Self {
+    pub async fn init(
+        context: E,
+        cfg: Config,
+        mailbox: buffered::Mailbox<P, Shard<B, H>>,
+        block_codec_cfg: B::Cfg,
+    ) -> Self {
+        let cfg = |name: &str| prunable::Config {
+            partition: format!("shard-layer-{}-{name}-map", cfg.partition_prefix),
+            translator: TwoCap,
+            items_per_section: cfg.items_per_section,
+            compression: None,
+            codec_config: (),
+            buffer_pool: cfg.buffer_pool.clone(),
+            replay_buffer: cfg.replay_buffer,
+            write_buffer: cfg.write_buffer,
+        };
+        let commitment_map =
+            prunable::Archive::init(context.with_label("commitment-map"), cfg("commitment"))
+                .await
+                .unwrap_or_else(|_| panic!("Failed to initialize commitment archive"));
+        let digest_map = prunable::Archive::init(context.with_label("digest-map"), cfg("digest"))
+            .await
+            .unwrap_or_else(|_| panic!("Failed to initialize digest archive"));
+
         Self {
             mailbox,
-            block_codec_cfg: cfg,
-            commitment_map: HashMap::new(),
-            reconstruction_cache: HashMap::new(),
+            block_codec_cfg,
+            commitment_map,
+            digest_map,
         }
     }
 
@@ -81,7 +130,7 @@ where
     ) {
         for (peer, chunk) in chunks {
             let message = Shard::new(coding_commitment, config, chunk);
-            let _peers = self.broadcast(Recipients::One(peer), message).await;
+            let _peers = self.mailbox.broadcast(Recipients::One(peer), message).await;
         }
     }
 
@@ -92,25 +141,10 @@ where
         let available_shards = self.mailbox.get(None, commitment, None).await;
 
         for shard in available_shards {
-            let _peers = self.broadcast(Recipients::All, shard).await;
+            let _peers = self.mailbox.broadcast(Recipients::All, shard).await;
         }
     }
 
-    /// Attempts to fetch a cached reconstructed [Block] by its digest.
-    pub fn get(&mut self, digest: B::Digest) -> Option<B> {
-        self.reconstruction_cache.get(&digest).cloned()
-    }
-
-    /// Checks if the shard layer has the digest corresponding to a given coding commitment.
-    pub fn has_digest(&self, commitment: &B::Commitment) -> bool {
-        self.commitment_map.contains_key(commitment)
-    }
-
-    /// Gets the digest corresponding to a given coding commitment, if known.
-    pub fn get_digest(&self, commitment: &B::Commitment) -> Option<B::Digest> {
-        self.commitment_map.get(commitment).copied()
-    }
-
     /// Attempts to retrieve and reconstruct a [Block] by its coding commitment from a set of [Shard]s
     /// received from peers.
     ///
@@ -121,12 +155,12 @@ where
     pub async fn try_reconstruct(
         &mut self,
         commitment: B::Commitment,
-    ) -> Result<(), ReconstructionError> {
+    ) -> Result<Option<B>, ReconstructionError> {
         let available_chunks = self.mailbox.get(None, commitment, None).await;
 
         let Some((total, min)) = available_chunks.first().map(|c| c.config) else {
             // No chunks available.
-            return Ok(());
+            return Ok(None);
         };
         let coded_chunks = available_chunks
             .iter()
@@ -142,7 +176,7 @@ where
                 need = min,
                 "not enough chunks to reconstruct block",
             );
-            return Ok(());
+            return Ok(None);
         }
 
         // Attempt to recover the block from the available chunks. This process will also
@@ -152,13 +186,18 @@ where
         // Attempt to decode the block from the recovered data.
         let block = B::decode_cfg(&mut recovered.as_slice(), &self.block_codec_cfg)?;
 
-        // ---- DEBUG ----
+        self.commitment_map
+            .put(block.height(), commitment, block.digest())
+            .await
+            .expect("failed to put commitment");
+        self.digest_map
+            .put(block.height(), block.digest(), commitment)
+            .await
+            .expect("failed to put digest");
+
         info!(%commitment, ?block, "successfully reconstructed block");
-        self.commitment_map.insert(commitment, block.digest());
-        self.reconstruction_cache.insert(block.digest(), block);
-        // ----
 
-        Ok(())
+        Ok(Some(block))
     }
 
     /// Subscribes to a block by commitment with an externally prepared responder.
@@ -173,25 +212,46 @@ where
     ) -> Result<(), ReconstructionError> {
         todo!("Subscribe to all chunks, reconstruct block when enough are available.");
     }
-}
-
-impl<P, B, H> Broadcaster for ShardLayer<P, B, H>
-where
-    P: PublicKey,
-    B: Block<Digest = H::Digest, Commitment = H::Digest>,
-    H: Hasher,
-{
-    type Recipients = Recipients<P>;
-    type Message = Shard<B, H>;
-    type Response = Vec<P>;
 
-    async fn broadcast(
+    /// Attempts to fetch a reconstructed [Block] by its digest.
+    pub async fn block_by_digest(
         &mut self,
-        recipients: Self::Recipients,
-        message: Self::Message,
-    ) -> oneshot::Receiver<Self::Response> {
-        // Direct broadcasts of individual chunks to the underlying mailbox.
-        self.mailbox.broadcast(recipients, message).await
+        digest: B::Digest,
+    ) -> Result<Option<B>, ReconstructionError> {
+        match self
+            .digest_map
+            .get(Identifier::Key(&digest))
+            .await
+            .expect("failed to get digest")
+        {
+            Some(commitment) => self.try_reconstruct(commitment).await,
+            None => Ok(None),
+        }
+    }
+
+    /// Checks if the shard layer has the digest corresponding to a given coding commitment.
+    pub async fn has_digest(&self, commitment: &B::Commitment) -> bool {
+        self.commitment_map
+            .has(Identifier::Key(commitment))
+            .await
+            .expect("failed to get commitment")
+    }
+
+    /// Gets the digest corresponding to a given coding commitment, if known.
+    pub async fn get_digest(&self, commitment: &B::Commitment) -> Option<B::Digest> {
+        self.commitment_map
+            .get(Identifier::Key(commitment))
+            .await
+            .expect("failed to get commitment")
+    }
+
+    /// Prunes old entries from the internal maps.
+    pub async fn prune(&mut self, up_to: u64) {
+        futures::try_join!(
+            self.commitment_map.prune(up_to),
+            self.digest_map.prune(up_to)
+        )
+        .expect("failed to prune maps");
     }
 }
 
diff --git a/consensus/src/marshal/mod.rs b/consensus/src/marshal/mod.rs
@@ -76,7 +76,7 @@ mod tests {
         resolver::p2p as resolver,
     };
     use crate::{
-        marshal::ingress::coding::ShardLayer,
+        marshal::ingress::coding::{self, ShardLayer},
         threshold_simplex::types::{
             finalize_namespace, notarize_namespace, seed_namespace, Activity, Finalization,
             Finalize, Notarization, Notarize, Proposal,
@@ -113,7 +113,7 @@ mod tests {
     use rand::{seq::SliceRandom, Rng};
     use std::{
         collections::BTreeMap,
-        num::{NonZeroU32, NonZeroUsize},
+        num::{NonZero, NonZeroU32, NonZeroUsize},
         time::Duration,
     };
 
@@ -200,10 +200,18 @@ mod tests {
             codec_config: (),
         };
         let (broadcast_engine, buffer) = buffered::Engine::new(context.clone(), broadcast_config);
-        let shards = ShardLayer::new(buffer, ());
         let network = oracle.register(secret.public_key(), 2).await.unwrap();
         broadcast_engine.start(network);
 
+        let shard_config = coding::Config {
+            partition_prefix: "shards".to_string(),
+            items_per_section: NonZero::new(100).unwrap(),
+            replay_buffer: NonZero::new(100).unwrap(),
+            write_buffer: NonZero::new(100).unwrap(),
+            buffer_pool: PoolRef::new(PAGE_SIZE, PAGE_CACHE_SIZE),
+        };
+        let shards = ShardLayer::init(context.with_label("shard"), shard_config, buffer, ()).await;
+
         let (actor, mailbox) = actor::Actor::init(context.clone(), config).await;
         let application = Application::<B>::default();
 
