perf: optimize DataMap::get_batched

Author: zvolin

node/store/src/rocksdb/map.rs

@@ -16,6 +16,8 @@
 
 use super::*;
 
+use indexmap::IndexMap;
+use rocksdb::WriteBatch;
 use snarkvm::synthesizer::store::helpers::{Map, MapRead};
 
 use core::{fmt, fmt::Debug, hash::Hash};
@@ -25,7 +27,10 @@ use std::{borrow::Cow, sync::atomic::Ordering};
 pub struct DataMap<K: Serialize + DeserializeOwned, V: Serialize + DeserializeOwned> {
     pub(super) database: RocksDB,
     pub(super) context: Vec<u8>,
-    pub(super) _phantom: PhantomData<(K, V)>,
+    /// The tracker for whether a database transaction is in progress.
+    pub(super) batch_in_progress: Arc<AtomicBool>,
+    /// The database transaction.
+    pub(super) atomic_batch: Arc<Mutex<IndexMap<K, Option<V>>>>,
 }
 
 impl<
@@ -38,18 +43,19 @@ impl<
     /// Inserts the given key-value pair into the map.
     ///
     fn insert(&self, key: K, value: V) -> Result<()> {
-        // Prepare the prefixed key and serialized value.
-        let raw_key = self.create_prefixed_key(&key)?;
-        let raw_value = bincode::serialize(&value)?;
-
         // Determine if an atomic batch is in progress.
-        let is_batch = self.database.batch_in_progress.load(Ordering::SeqCst);
+        let is_batch = self.batch_in_progress.load(Ordering::SeqCst);
 
         match is_batch {
             // If a batch is in progress, add the key-value pair to the batch.
-            true => self.database.atomic_batch.lock().put(raw_key, raw_value),
+            true => {
+                self.atomic_batch.lock().insert(key, Some(value));
+            }
             // Otherwise, insert the key-value pair directly into the map.
             false => {
+                // Prepare the prefixed key and serialized value.
+                let raw_key = self.create_prefixed_key(&key)?;
+                let raw_value = bincode::serialize(&value)?;
                 self.database.put(&raw_key, &raw_value)?;
             }
         }
@@ -61,17 +67,18 @@ impl<
     /// Removes the key-value pair for the given key from the map.
     ///
     fn remove(&self, key: &K) -> Result<()> {
-        // Prepare the prefixed key.
-        let raw_key = self.create_prefixed_key(key)?;
-
         // Determine if an atomic batch is in progress.
-        let is_batch = self.database.batch_in_progress.load(Ordering::SeqCst);
+        let is_batch = self.batch_in_progress.load(Ordering::SeqCst);
 
         match is_batch {
             // If a batch is in progress, add the key to the batch.
-            true => self.database.atomic_batch.lock().delete(raw_key),
+            true => {
+                self.atomic_batch.lock().insert(*key, None);
+            }
             // Otherwise, remove the key-value pair directly from the map.
             false => {
+                // Prepare the prefixed key.
+                let raw_key = self.create_prefixed_key(key)?;
                 self.database.delete(&raw_key)?;
             }
         }
@@ -85,9 +92,9 @@ impl<
     ///
     fn start_atomic(&self) {
         // Set the atomic batch flag to `true`.
-        self.database.batch_in_progress.store(true, Ordering::SeqCst);
+        self.batch_in_progress.store(true, Ordering::SeqCst);
         // Ensure that the atomic batch is empty.
-        assert!(self.database.atomic_batch.lock().is_empty());
+        assert!(self.atomic_batch.lock().is_empty());
     }
 
     ///
@@ -96,33 +103,50 @@ impl<
     /// if they are already part of a larger one.
     ///
     fn is_atomic_in_progress(&self) -> bool {
-        self.database.batch_in_progress.load(Ordering::SeqCst)
+        self.batch_in_progress.load(Ordering::SeqCst)
     }
 
     ///
     /// Aborts the current atomic operation.
     ///
     fn abort_atomic(&self) {
         // Clear the atomic batch.
-        self.database.atomic_batch.lock().clear();
+        *self.atomic_batch.lock() = Default::default();
         // Set the atomic batch flag to `false`.
-        self.database.batch_in_progress.store(false, Ordering::SeqCst);
+        self.batch_in_progress.store(false, Ordering::SeqCst);
     }
 
     ///
     /// Finishes an atomic operation, performing all the queued writes.
     ///
     fn finish_atomic(&self) -> Result<()> {
         // Retrieve the atomic batch.
-        let batch = core::mem::take(&mut *self.database.atomic_batch.lock());
-
-        if !batch.is_empty() {
-            // Execute the batch of operations atomically.
+        let operations = core::mem::take(&mut *self.atomic_batch.lock());
+
+        if !operations.is_empty() {
+            // Prepare operations batch for underlying database.
+            let mut batch = WriteBatch::default();
+            for operation in operations {
+                match operation {
+                    (key, Some(value)) => {
+                        // Prepare the prefixed key and serialized value for insertion.
+                        let raw_key = self.create_prefixed_key(&key)?;
+                        let raw_value = bincode::serialize(&value)?;
+                        batch.put(raw_key, raw_value);
+                    }
+                    (key, None) => {
+                        // Prepare the prefixed key for deletion.
+                        let raw_key = self.create_prefixed_key(&key)?;
+                        batch.delete(raw_key);
+                    }
+                };
+            }
+            // Execute all the operations atomically.
             self.database.rocksdb.write(batch)?;
         }
 
         // Set the atomic batch flag to `false`.
-        self.database.batch_in_progress.store(false, Ordering::SeqCst);
+        self.batch_in_progress.store(false, Ordering::SeqCst);
 
         Ok(())
     }
@@ -177,60 +201,7 @@ impl<
         K: Borrow<Q>,
         Q: PartialEq + Eq + Hash + Serialize + ?Sized,
     {
-        // Return early if there is no atomic batch in progress.
-        if self.database.batch_in_progress.load(Ordering::SeqCst) {
-            struct OperationFinder {
-                key: Vec<u8>,
-                value: Option<Option<Box<[u8]>>>,
-            }
-
-            impl rocksdb::WriteBatchIterator for OperationFinder {
-                fn put(&mut self, key: Box<[u8]>, value: Box<[u8]>) {
-                    if *key == self.key {
-                        self.value = Some(Some(value));
-                    }
-                }
-
-                fn delete(&mut self, key: Box<[u8]>) {
-                    if *key == self.key {
-                        self.value = Some(None);
-                    }
-                }
-            }
-
-            // Prepare the prefixed key and serialized value.
-            let raw_key = match self.create_prefixed_key(key) {
-                Ok(key) => key,
-                Err(error) => {
-                    error!("Failed to create prefixed key in 'get_batched': {:?}", error);
-                    return None;
-                }
-            };
-
-            // Retrieve the atomic batch.
-            let batch = self.database.atomic_batch.lock();
-
-            // Initialize the operation finder.
-            let mut finder = OperationFinder { key: raw_key, value: None };
-
-            // Iterate over the batch.
-            batch.iterate(&mut finder);
-
-            // Return the value.
-            match finder.value {
-                Some(Some(value)) => match bincode::deserialize(&value) {
-                    Ok(value) => Some(Some(value)),
-                    Err(error) => {
-                        error!("Failed to deserialize value in 'get_batched': {:?}", error);
-                        None
-                    }
-                },
-                Some(None) => Some(None),
-                None => None,
-            }
-        } else {
-            None
-        }
+        if self.batch_in_progress.load(Ordering::SeqCst) { self.atomic_batch.lock().get(key).cloned() } else { None }
     }
 
     ///

node/store/src/rocksdb/mod.rs

@@ -62,10 +62,6 @@ pub struct RocksDB {
     network_id: u16,
     /// The optional development ID.
     dev: Option<u16>,
-    /// The tracker for whether a database transaction is in progress.
-    batch_in_progress: Arc<AtomicBool>,
-    /// The database transaction.
-    atomic_batch: Arc<Mutex<rocksdb::WriteBatch>>,
 }
 
 impl Deref for RocksDB {
@@ -102,13 +98,7 @@ impl Database for RocksDB {
                     Arc::new(rocksdb::DB::open(&options, primary)?)
                 };
 
-                Ok::<_, anyhow::Error>(RocksDB {
-                    rocksdb,
-                    network_id,
-                    dev,
-                    batch_in_progress: Default::default(),
-                    atomic_batch: Default::default(),
-                })
+                Ok::<_, anyhow::Error>(RocksDB { rocksdb, network_id, dev })
             })?
             .clone();
 
@@ -133,7 +123,7 @@ impl Database for RocksDB {
         context.extend_from_slice(&(data_id as u16).to_le_bytes());
 
         // Return the DataMap.
-        Ok(DataMap { database, context, _phantom: PhantomData })
+        Ok(DataMap { database, context, batch_in_progress: Default::default(), atomic_batch: Default::default() })
     }
 }
 
@@ -162,13 +152,7 @@ impl RocksDB {
                 Arc::new(rocksdb::DB::open(&options, primary)?)
             };
 
-            Ok::<_, anyhow::Error>(RocksDB {
-                rocksdb,
-                network_id: u16::MAX,
-                dev,
-                batch_in_progress: Default::default(),
-                atomic_batch: Default::default(),
-            })
+            Ok::<_, anyhow::Error>(RocksDB { rocksdb, network_id: u16::MAX, dev })
         }?;
 
         // Ensure the database development ID match.
@@ -193,7 +177,7 @@ impl RocksDB {
         context.extend_from_slice(&(data_id as u16).to_le_bytes());
 
         // Return the DataMap.
-        Ok(DataMap { database, context, _phantom: PhantomData })
+        Ok(DataMap { database, context, batch_in_progress: Default::default(), atomic_batch: Default::default() })
     }
 }