[storage/journal] batch read/append journal optimizations (#3709)

Co-authored-by: Dan Laine <dan@commonware.xyz>

Author: roberto-bayardo

storage/src/journal/authenticated.rs

@@ -27,7 +27,7 @@ use commonware_codec::{CodecFixedShared, CodecShared, Encode, EncodeShared};
 use commonware_cryptography::{Digest, Hasher};
 use commonware_parallel::{Sequential, Strategy};
 use core::num::NonZeroU64;
-use futures::{future::try_join_all, try_join, TryFutureExt as _};
+use futures::{try_join, TryFutureExt as _};
 use thiserror::Error;
 use tracing::{debug, warn};
 
@@ -646,10 +646,8 @@ where
             )
             .await?;
 
-        let futures = (*start_loc..*end_loc)
-            .map(|i| reader.read(i))
-            .collect::<Vec<_>>();
-        let ops = try_join_all(futures).await?;
+        let positions: Vec<u64> = (*start_loc..*end_loc).collect();
+        let ops = reader.read_many(&positions).await?;
 
         Ok((proof, ops))
     }

storage/src/journal/benches/bench.rs

@@ -19,13 +19,15 @@ mod fixed_append;
 mod fixed_read_random;
 mod fixed_read_sequential;
 mod fixed_replay;
+mod variable_read_random;
 mod variable_replay;
 
 criterion_main!(
     fixed_append::benches,
     fixed_read_random::benches,
     fixed_read_sequential::benches,
     fixed_replay::benches,
+    variable_read_random::benches,
     variable_replay::benches,
 );
 

storage/src/journal/benches/variable_read_random.rs

@@ -0,0 +1,142 @@
+use crate::{append_fixed_random_data, get_variable_journal};
+use commonware_runtime::{
+    benchmarks::{context, tokio},
+    tokio::{Config, Context, Runner},
+    Runner as _, Supervisor as _,
+};
+use commonware_storage::journal::contiguous::{variable::Journal, Reader as _};
+use commonware_utils::{sequence::FixedBytes, NZU64};
+use criterion::{criterion_group, Criterion};
+use futures::future::try_join_all;
+use rand::{rngs::StdRng, Rng, SeedableRng};
+use std::{
+    hint::black_box,
+    num::NonZeroU64,
+    time::{Duration, Instant},
+};
+
+/// Partition name to use in the journal config.
+const PARTITION: &str = "variable-random-test-partition";
+
+/// Value of items_per_section to use in the journal config.
+const ITEMS_PER_SECTION: NonZeroU64 = NZU64!(10_000);
+
+/// Number of items to write to the journal we will be reading from.
+const ITEMS_TO_WRITE: u64 = 5_000_000;
+
+/// Size of each journal item in bytes.
+const ITEM_SIZE: usize = 32;
+
+/// Read `items_to_read` random items from the given `journal`, awaiting each
+/// result before continuing.
+async fn bench_run_serial(journal: &Journal<Context, FixedBytes<ITEM_SIZE>>, items_to_read: usize) {
+    let reader = journal.reader().await;
+    let mut rng = StdRng::seed_from_u64(0);
+    for _ in 0..items_to_read {
+        let pos = rng.gen_range(0..ITEMS_TO_WRITE);
+        black_box(reader.read(pos).await.expect("failed to read data"));
+    }
+}
+
+/// Concurrently read (via try_join_all) `items_to_read` random items from the given `journal`.
+async fn bench_run_concurrent(
+    journal: &Journal<Context, FixedBytes<ITEM_SIZE>>,
+    items_to_read: usize,
+) {
+    let reader = journal.reader().await;
+    let mut rng = StdRng::seed_from_u64(0);
+    let mut futures = Vec::with_capacity(items_to_read);
+    for _ in 0..items_to_read {
+        let pos = rng.gen_range(0..ITEMS_TO_WRITE);
+        futures.push(reader.read(pos));
+    }
+    try_join_all(futures).await.expect("failed to read data");
+}
+
+/// Batch-read `items_to_read` random items via `read_many`.
+async fn bench_run_read_many(
+    journal: &Journal<Context, FixedBytes<ITEM_SIZE>>,
+    items_to_read: usize,
+) {
+    let reader = journal.reader().await;
+    let mut rng = StdRng::seed_from_u64(0);
+    let mut positions: Vec<u64> = (0..items_to_read)
+        .map(|_| rng.gen_range(0..ITEMS_TO_WRITE))
+        .collect();
+    positions.sort_unstable();
+    positions.dedup();
+    black_box(
+        reader
+            .read_many(&positions)
+            .await
+            .expect("failed to read data"),
+    );
+}
+
+fn bench_variable_read_random(c: &mut Criterion) {
+    let cfg = Config::default();
+    let mut initialized = false;
+    let runner = tokio::Runner::new(cfg.clone());
+    for mode in ["serial", "concurrent", "read_many"] {
+        for items_to_read in [100, 1_000, 10_000, 100_000] {
+            c.bench_function(
+                &format!(
+                    "{}/mode={} items={} size={}",
+                    module_path!(),
+                    mode,
+                    items_to_read,
+                    ITEM_SIZE
+                ),
+                |b| {
+                    // Setup: populate journal (once, on first sample).
+                    if !initialized {
+                        Runner::new(cfg.clone()).start(|ctx| async move {
+                            let mut j =
+                                get_variable_journal(ctx, PARTITION, ITEMS_PER_SECTION).await;
+                            append_fixed_random_data::<_, ITEM_SIZE>(&mut j, ITEMS_TO_WRITE).await;
+                            j.sync().await.unwrap();
+                        });
+                        initialized = true;
+                    }
+
+                    // Benchmark: measure read time.
+                    b.to_async(&runner).iter_custom(|iters| async move {
+                        let ctx = context::get::<commonware_runtime::tokio::Context>();
+                        let j = get_variable_journal(
+                            ctx.child("storage"),
+                            PARTITION,
+                            ITEMS_PER_SECTION,
+                        )
+                        .await;
+                        let mut duration = Duration::ZERO;
+                        for _ in 0..iters {
+                            let start = Instant::now();
+                            match mode {
+                                "serial" => bench_run_serial(&j, items_to_read).await,
+                                "concurrent" => bench_run_concurrent(&j, items_to_read).await,
+                                "read_many" => bench_run_read_many(&j, items_to_read).await,
+                                _ => unreachable!(),
+                            }
+                            duration += start.elapsed();
+                        }
+                        duration
+                    });
+                },
+            );
+        }
+    }
+
+    // Cleanup: destroy journal.
+    if initialized {
+        Runner::new(cfg).start(|context| async move {
+            let j = get_variable_journal::<ITEM_SIZE>(context, PARTITION, ITEMS_PER_SECTION).await;
+            j.destroy().await.unwrap();
+        });
+    }
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default().sample_size(10);
+    targets = bench_variable_read_random
+}

storage/src/journal/contiguous/fixed.rs

@@ -158,14 +158,20 @@ impl<E: Context, A: CodecFixedShared> Inner<E, A> {
 
     /// Read an item if it can be done synchronously (e.g. without I/O), returning `None` otherwise.
     fn try_read_sync(&self, pos: u64, items_per_blob: u64) -> Option<A> {
+        let mut buf = vec![0u8; SegmentedJournal::<E, A>::CHUNK_SIZE];
+        self.try_read_sync_into(pos, items_per_blob, &mut buf)
+    }
+
+    /// Read an item synchronously using caller-provided buffer.
+    fn try_read_sync_into(&self, pos: u64, items_per_blob: u64, buf: &mut [u8]) -> Option<A> {
         if pos >= self.size || pos < self.pruning_boundary {
             return None;
         }
         let section = pos / items_per_blob;
         let section_start = section * items_per_blob;
         let first_in_section = self.pruning_boundary.max(section_start);
         let pos_in_section = pos - first_in_section;
-        self.journal.try_get_sync(section, pos_in_section)
+        self.journal.try_get_sync_into(section, pos_in_section, buf)
     }
 }
 
@@ -238,8 +244,12 @@ impl<E: Context, A: CodecFixedShared> super::Reader for Reader<'_, E, A> {
         let mut miss_indices: Vec<usize> = Vec::new();
         let mut miss_positions: Vec<u64> = Vec::new();
 
+        let mut sync_buf = vec![0u8; chunk_size];
         for (i, &pos) in positions.iter().enumerate() {
-            if let Some(item) = self.guard.try_read_sync(pos, items_per_blob) {
+            if let Some(item) = self
+                .guard
+                .try_read_sync_into(pos, items_per_blob, &mut sync_buf)
+            {
                 result.push(Some(item));
             } else {
                 result.push(None);

storage/src/journal/contiguous/mod.rs

@@ -35,11 +35,11 @@ pub trait Reader: Send + Sync {
     /// Guaranteed not to return [Error::ItemPruned] for positions within `bounds()`.
     fn read(&self, position: u64) -> impl Future<Output = Result<Self::Item, Error>> + Send;
 
-    /// Read multiple items at the given positions, which must be sorted in ascending order.
+    /// Read multiple items at the given positions. `positions` must be sorted in strictly ascending
+    /// order (sorted and unique).
     ///
-    /// The default implementation calls [`read`](Self::read) in a loop.
-    /// Fixed-size journal implementations override this to amortize lock
-    /// acquisition and avoid per-item buffer allocation.
+    /// The default implementation calls [`read`](Self::read) in a loop. Concrete journal
+    /// implementations override this to amortize lock acquisition and batch I/O.
     fn read_many(
         &self,
         positions: &[u64],
@@ -111,6 +111,14 @@ pub enum Many<'a, T> {
 }
 
 impl<T> Many<'_, T> {
+    /// Returns the total number of items across all segments.
+    pub fn len(&self) -> usize {
+        match self {
+            Self::Flat(items) => items.len(),
+            Self::Nested(nested_items) => nested_items.iter().map(|items| items.len()).sum(),
+        }
+    }
+
     /// Returns `true` if there are no items across all segments.
     pub fn is_empty(&self) -> bool {
         match self {

storage/src/journal/contiguous/tests.rs

@@ -68,6 +68,7 @@ where
     test_persistence_basic(&indexed_factory).await;
     test_persistence_after_prune(&indexed_factory).await;
     test_read_by_position(&indexed_factory).await;
+    test_read_many(&indexed_factory).await;
     test_read_out_of_range(&indexed_factory).await;
     test_read_after_prune(&indexed_factory).await;
     test_rewind_to_middle(&indexed_factory).await;
@@ -732,6 +733,26 @@ where
     journal.destroy().await.unwrap();
 }
 
+/// Test reading multiple items by position.
+pub(super) async fn test_read_many<F, J>(factory: &F)
+where
+    F: Fn(String) -> BoxFuture<'static, Result<J, Error>>,
+    J: PersistableContiguous,
+{
+    let mut journal = factory("read-many".into()).await.unwrap();
+
+    for i in 0..15u64 {
+        journal.append(&(i * 100)).await.unwrap();
+    }
+
+    let reader = journal.reader().await;
+    let items = reader.read_many(&[1, 4, 12]).await.unwrap();
+    assert_eq!(items, vec![100, 400, 1200]);
+    drop(reader);
+
+    journal.destroy().await.unwrap();
+}
+
 /// Test read errors for out-of-range positions.
 pub(super) async fn test_read_out_of_range<F, J>(factory: &F)
 where