session: hash actuals for the merkle update on multiple threads

nomt/src/lib.rs

@@ -197,6 +197,7 @@ impl From<[u8; 32]> for Root {
 /// An instance of the Nearly-Optimal Merkle Trie Database.
 pub struct Nomt<T> {
     merkle_update_pool: UpdatePool,
+    actuals_concurrency: usize,
     /// The handle to the page cache.
     page_cache: PageCache,
     page_pool: PagePool,
@@ -237,6 +238,7 @@ impl<T: HashAlgorithm> Nomt<T> {
 
         Ok(Self {
             merkle_update_pool: UpdatePool::new(o.commit_concurrency, o.warm_up),
+            actuals_concurrency: o.actuals_concurrency,
             page_cache,
             page_pool,
             store,
@@ -328,6 +330,7 @@ impl<T: HashAlgorithm> Nomt<T> {
 
         Session {
             store,
+            actuals_concurrency: self.actuals_concurrency,
             merkle_updater: self.merkle_update_pool.begin::<T>(
                 self.page_cache.clone(),
                 self.page_pool.clone(),
@@ -477,6 +480,7 @@ impl SessionParams {
 /// correctness of replaying the same operations.
 pub struct Session<T> {
     store: Store,
+    actuals_concurrency: usize,
     merkle_updater: Updater,
     metrics: Metrics,
     rollback_delta: Option<rollback::ReverseDeltaBuilder>,
@@ -585,14 +589,10 @@ impl<T: HashAlgorithm> Session<T> {
             .take()
             .map(|delta_builder| delta_builder.finalize(&actuals));
 
-        let mut compact_actuals = Vec::with_capacity(actuals.len());
-        for (path, read_write) in &actuals {
-            compact_actuals.push((path.clone(), read_write.to_compact::<T>()));
-        }
-
-        let merkle_update_handle = self
-            .merkle_updater
-            .update_and_prove::<T>(compact_actuals, self.witness_mode.0)?;
+        let merkle_update_handle = self.merkle_updater.update_and_prove::<T>(
+            compact_actuals::<T>(&actuals, self.actuals_concurrency),
+            self.witness_mode.0,
+        )?;
 
         let mut tx = self.store.new_value_tx();
         for (path, read_write) in actuals {
@@ -613,6 +613,52 @@ impl<T: HashAlgorithm> Session<T> {
     }
 }
 
+/// Hashes the values of `actuals` for the merkle update, splitting large inputs across scoped threads.
+fn compact_actuals<T: HashAlgorithm>(
+    actuals: &[(KeyPath, KeyReadWrite)],
+    max_threads: usize,
+) -> Vec<(KeyPath, crate::merkle::KeyReadWrite)> {
+    const MIN_HASHES_PER_THREAD: usize = 8192;
+    let compact = |chunk: &[(KeyPath, KeyReadWrite)]| {
+        chunk
+            .iter()
+            .map(|(path, read_write)| (*path, read_write.to_compact::<T>()))
+            .collect::<Vec<_>>()
+    };
+    if max_threads == 1 || actuals.len() < 2 * MIN_HASHES_PER_THREAD {
+        return compact(actuals);
+    }
+    let hashes = actuals
+        .iter()
+        .filter(|(_, read_write)| {
+            matches!(
+                read_write,
+                KeyReadWrite::Write(Some(_)) | KeyReadWrite::ReadThenWrite(_, Some(_))
+            )
+        })
+        .count();
+    let threads = max_threads.min(hashes / MIN_HASHES_PER_THREAD);
+    if threads < 2 {
+        return compact(actuals);
+    }
+    let threads = std::thread::available_parallelism().map_or(1, |n| n.get().min(threads));
+    if threads == 1 {
+        return compact(actuals);
+    }
+    let chunk = actuals.len().div_ceil(threads);
+    std::thread::scope(|scope| {
+        let handles: Vec<_> = actuals
+            .chunks(chunk)
+            .map(|chunk| scope.spawn(move || compact(chunk)))
+            .collect();
+        let mut output = Vec::with_capacity(actuals.len());
+        for handle in handles {
+            output.extend(handle.join().expect("compacting thread panicked"));
+        }
+        output
+    })
+}
+
 /// A finished session.
 ///
 /// This is the result of completing a session and computing the merkle root and merkle DB changes,
@@ -953,6 +999,48 @@ pub fn validate_hashtable(options: &Options) -> anyhow::Result<HashTableUtilizat
 mod tests {
     use crate::hasher::Blake3Hasher;
 
+    #[test]
+    fn parallel_compaction_matches_sequential() {
+        use super::{compact_actuals, KeyReadWrite};
+        use crate::merkle::KeyReadWrite as Compact;
+
+        let actuals: Vec<_> = (0u64..131_073)
+            .map(|i| {
+                let mut path = [0; 32];
+                path[..8].copy_from_slice(&i.to_be_bytes());
+                let value = Some(i.to_le_bytes().to_vec());
+                let read_write = match i % 6 {
+                    0 => KeyReadWrite::Read(None),
+                    1 => KeyReadWrite::Read(value),
+                    2 => KeyReadWrite::Write(None),
+                    3 => KeyReadWrite::Write(value),
+                    4 => KeyReadWrite::ReadThenWrite(value, None),
+                    _ => KeyReadWrite::ReadThenWrite(None, value),
+                };
+                (path, read_write)
+            })
+            .collect();
+
+        for len in [0, 1, 16_383, 16_384, 49_151, 49_152, 49_153, 131_073] {
+            let sequential = compact_actuals::<Blake3Hasher>(&actuals[..len], 1);
+            for threads in [2, 16] {
+                let parallel = compact_actuals::<Blake3Hasher>(&actuals[..len], threads);
+                assert_eq!(parallel.len(), sequential.len());
+                for ((path, op), (expected_path, expected_op)) in parallel.iter().zip(&sequential) {
+                    assert_eq!(path, expected_path);
+                    match (op, expected_op) {
+                        (Compact::Read, Compact::Read) => {}
+                        (Compact::Write(value), Compact::Write(expected))
+                        | (Compact::ReadThenWrite(value), Compact::ReadThenWrite(expected)) => {
+                            assert_eq!(value, expected);
+                        }
+                        _ => panic!("operation changed: {op:?} != {expected_op:?}"),
+                    }
+                }
+            }
+        }
+    }
+
     #[test]
     fn session_is_sync() {
         fn is_sync<T: Sync>() {}

nomt/src/options.rs

@@ -7,6 +7,7 @@ pub struct Options {
     pub(crate) path: PathBuf,
     /// The number of commit workers. Values over 64 will be rounded down to 64.
     pub(crate) commit_concurrency: usize,
+    pub(crate) actuals_concurrency: usize,
     /// The number of io_uring instances, or I/O threads on non-Linux platforms.
     pub(crate) io_workers: usize,
     /// Enable or disable metrics collection.
@@ -40,6 +41,7 @@ impl Options {
         Self {
             path: PathBuf::from("nomt_db"),
             commit_concurrency: 1,
+            actuals_concurrency: 16,
             io_workers: 3,
             metrics: false,
             bitbox_num_pages: 64_000,
@@ -73,6 +75,19 @@ impl Options {
         self.commit_concurrency = commit_concurrency;
     }
 
+    /// Set the maximum number of threads hashing actuals per session.
+    ///
+    /// Limited by available CPU parallelism. Set to 1 to hash sequentially.
+    /// Independent of commit and I/O worker concurrency.
+    ///
+    /// Default: 16.
+    ///
+    /// Panics if `max_threads` is zero.
+    pub fn actuals_concurrency(&mut self, max_threads: usize) {
+        assert!(max_threads > 0);
+        self.actuals_concurrency = max_threads;
+    }
+
     /// Set metrics collection on or off.
     ///
     /// Default: off.