Refactor to avoid holding database transactions longer than necessary

Author: Tumas

database/src/lib.rs

@@ -184,6 +184,39 @@ impl Database {
         Ok(())
     }
 
+    pub fn delete_batch(&self, keys: impl IntoIterator<Item = impl AsRef<[u8]>>) -> Result<()> {
+        match self.kind() {
+            #[cfg(not(target_os = "zkvm"))]
+            DatabaseKind::Persistent {
+                database_name,
+                environment,
+                restart_tx: _,
+            } => {
+                let transaction = environment.begin_rw_txn()?;
+                let database = transaction.open_db(Some(database_name))?;
+
+                let mut cursor = transaction.cursor(database.dbi())?;
+
+                for key in keys {
+                    if cursor.set::<()>(key.as_ref())?.is_some() {
+                        cursor.del(WriteFlags::default())?;
+                    }
+                }
+
+                transaction.commit()?;
+            }
+            DatabaseKind::InMemory { map } => {
+                let mut map = map.lock().expect("in-memory database mutex is poisoned");
+
+                for key in keys {
+                    map.remove(key.as_ref());
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     pub fn delete_range(&self, range: Range<impl AsRef<[u8]>>) -> Result<()> {
         let start = range.start.as_ref();
         let end = range.end.as_ref();
@@ -436,7 +469,6 @@ impl Database {
             } => {
                 let transaction = environment.begin_ro_txn()?;
                 let database = transaction.open_db(Some(database_name))?;
-
                 let mut cursor = transaction.cursor(database.dbi())?;
 
                 let first = if let Some((is_next, key, value)) = cursor.set_lowerbound(end)? {
@@ -500,14 +532,22 @@ impl Database {
                 environment,
                 restart_tx,
             } => {
+                let compressed_pairs = pairs
+                    .into_iter()
+                    .map(|(key, value)| Ok((key, compress(value.as_ref())?)))
+                    .collect::<Result<Vec<_>>>()?;
+
                 let transaction = environment.begin_rw_txn()?;
                 let database = transaction.open_db(Some(database_name))?;
 
-                for (key, value) in pairs {
-                    let key = key.as_ref();
-                    let compressed = compress(value.as_ref())?;
+                for (key, compressed) in &compressed_pairs {
                     transaction
-                        .put(database.dbi(), key, compressed, WriteFlags::default())
+                        .put(
+                            database.dbi(),
+                            key.as_ref(),
+                            compressed,
+                            WriteFlags::default(),
+                        )
                         .map_err(|error| {
                             handle_write_error(database_name, error, restart_tx.as_ref())
                         })?;
@@ -709,6 +749,21 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(build_persistent_database)]
+    #[test_case(build_in_memory_database)]
+    fn test_delete_batch(constructor: Constructor) -> Result<()> {
+        let database = constructor()?;
+
+        database.delete_batch(["A", "C", "D"])?;
+
+        assert_pairs_eq(
+            database.iterator_ascending("A"..)?,
+            [("B", "2"), ("E", "5")],
+        )?;
+
+        Ok(())
+    }
+
     #[test_case(build_persistent_database)]
     #[test_case(build_in_memory_database)]
     fn test_delete_range_inclusive_exclusive(constructor: Constructor) -> Result<()> {

fork_choice_control/src/storage.rs

@@ -409,72 +409,47 @@ impl<P: Preset> Storage<P> {
     }
 
     pub(crate) fn prune_old_blob_sidecars(&self, up_to_slot: Slot) -> Result<()> {
-        let mut blobs_to_remove: Vec<BlobIdentifier> = vec![];
-        let mut keys_to_remove = vec![];
-
         let results = self
             .database
             .iterator_descending(..=SlotBlobId(up_to_slot, H256::zero(), 0).to_string())?;
 
-        for result in results {
-            let (key_bytes, value_bytes) = result?;
+        let (mut keys_to_remove, blobs_to_remove): (Vec<_>, Vec<_>) =
+            itertools::process_results(results, |iter| {
+                iter.take_while(|(key_bytes, _)| SlotBlobId::has_prefix(key_bytes))
+                    .map(|(k, v)| (k.into_owned(), v))
+                    .unzip()
+            })?;
 
-            if !SlotBlobId::has_prefix(&key_bytes) {
-                break;
-            }
+        for blob_bytes in blobs_to_remove {
+            let BlobIdentifier { block_root, index } =
+                BlobIdentifier::from_ssz_default(blob_bytes)?;
 
-            // Deserialize-serialize BlobIdentifier as an additional measure
-            // to prevent other types of data getting accidentally deleted.
-            blobs_to_remove.push(BlobIdentifier::from_ssz_default(value_bytes)?);
-            keys_to_remove.push(key_bytes.into_owned());
+            keys_to_remove.push(BlobSidecarByBlobId(block_root, index).to_string().into());
         }
 
-        for blob_id in blobs_to_remove {
-            let BlobIdentifier { block_root, index } = blob_id;
-            let key = BlobSidecarByBlobId(block_root, index).to_string();
-
-            self.database.delete(key)?;
-        }
-
-        for key in keys_to_remove {
-            self.database.delete(key)?;
-        }
-
-        Ok(())
+        self.database.delete_batch(keys_to_remove)
     }
 
     pub(crate) fn prune_old_blocks_and_states(&self, up_to_slot: Slot) -> Result<()> {
-        let mut block_roots_to_remove = vec![];
-        let mut keys_to_remove = vec![];
-
         let results = self
             .database
             .iterator_descending(..=BlockRootBySlot(up_to_slot.saturating_sub(1)).to_string())?;
 
-        for result in results {
-            let (key_bytes, value_bytes) = result?;
-
-            if !BlockRootBySlot::has_prefix(&key_bytes) {
-                break;
-            }
-
-            block_roots_to_remove.push(H256::from_ssz_default(value_bytes)?);
-            keys_to_remove.push(key_bytes.into_owned());
-        }
-
-        for block_root in block_roots_to_remove {
-            let key = FinalizedBlockByRoot(block_root).to_string();
-            self.database.delete(key)?;
+        let (mut keys_to_remove, block_roots_to_remove): (Vec<_>, Vec<_>) =
+            itertools::process_results(results, |iter| {
+                iter.take_while(|(key_bytes, _)| BlockRootBySlot::has_prefix(key_bytes))
+                    .map(|(k, v)| (k.into_owned(), v))
+                    .unzip()
+            })?;
 
-            let key = StateByBlockRoot(block_root).to_string();
-            self.database.delete(key)?;
-        }
+        for block_root_bytes in block_roots_to_remove {
+            let block_root = H256::from_ssz_default(block_root_bytes)?;
 
-        for key in keys_to_remove {
-            self.database.delete(key)?;
+            keys_to_remove.push(FinalizedBlockByRoot(block_root).to_string().into());
+            keys_to_remove.push(StateByBlockRoot(block_root).to_string().into());
         }
 
-        Ok(())
+        self.database.delete_batch(keys_to_remove)
     }
 
     pub(crate) fn prune_old_state_roots(&self, up_to_slot: Slot) -> Result<()> {
@@ -484,25 +459,21 @@ impl<P: Preset> Storage<P> {
             .database
             .iterator_ascending(SlotByStateRoot(H256::zero()).to_string()..)?;
 
-        for result in results {
-            let (key_bytes, value_bytes) = result?;
-
-            if !SlotByStateRoot::has_prefix(&key_bytes) {
-                break;
-            }
+        let results = itertools::process_results(results, |iter| {
+            iter.take_while(|(key_bytes, _)| SlotByStateRoot::has_prefix(key_bytes))
+                .map(|(k, v)| (k.into_owned(), v))
+                .collect::<Vec<_>>()
+        })?;
 
+        for (key_bytes, value_bytes) in results {
             let slot = Slot::from_ssz_default(value_bytes)?;
 
             if slot < up_to_slot {
-                keys_to_remove.push(key_bytes.into_owned());
+                keys_to_remove.push(key_bytes);
             }
         }
 
-        for key in keys_to_remove {
-            self.database.delete(key)?;
-        }
-
-        Ok(())
+        self.database.delete_batch(keys_to_remove)
     }
 
     pub(crate) fn prune_unfinalized_blocks(&self, last_finalized_slot: Slot) -> Result<Vec<Slot>> {
@@ -513,19 +484,19 @@ impl<P: Preset> Storage<P> {
             .database
             .iterator_ascending(serialize_key(UnfinalizedBlockByRoot(H256::zero()))..)?;
 
-        for result in results {
-            let (key_bytes, value_bytes) = result?;
-
-            if !UnfinalizedBlockByRoot::has_prefix(&key_bytes) {
-                break;
-            }
+        let results = itertools::process_results(results, |iter| {
+            iter.take_while(|(key_bytes, _)| UnfinalizedBlockByRoot::has_prefix(key_bytes))
+                .map(|(k, v)| (k.into_owned(), v))
+                .collect::<Vec<_>>()
+        })?;
 
+        for (key_bytes, value_bytes) in results {
             let unfinalized_block = SignedBeaconBlock::<P>::from_ssz(&self.config, value_bytes)?;
             let block_slot = unfinalized_block.message().slot();
 
             if block_slot <= last_finalized_slot {
                 slots.push(block_slot);
-                keys_to_remove.push(key_bytes.into_owned());
+                keys_to_remove.push(key_bytes);
             }
         }
 
@@ -539,9 +510,7 @@ impl<P: Preset> Storage<P> {
             }
         }
 
-        for key in keys_to_remove {
-            self.database.delete(key)?;
-        }
+        self.database.delete_batch(keys_to_remove)?;
 
         Ok(slots)
     }
@@ -591,38 +560,29 @@ impl<P: Preset> Storage<P> {
     }
 
     pub(crate) fn prune_old_data_column_sidecars(&self, up_to_slot: Slot) -> Result<()> {
-        let mut columns_to_remove: Vec<DataColumnIdentifier> = vec![];
-        let mut keys_to_remove = vec![];
-
         let results = self
             .database
             .iterator_descending(..=SlotColumnId(up_to_slot, H256::zero(), 0).to_string())?;
 
-        for result in results {
-            let (key_bytes, value_bytes) = result?;
-
-            if !SlotColumnId::has_prefix(&key_bytes) {
-                break;
-            }
-
-            // Deserialize-serialize DataColumnIdentifier as an additional measure
-            // to prevent other types of data getting accidentally deleted.
-            columns_to_remove.push(DataColumnIdentifier::from_ssz_default(value_bytes)?);
-            keys_to_remove.push(key_bytes.into_owned());
+        let (mut keys_to_remove, columns_to_remove): (Vec<_>, Vec<_>) =
+            itertools::process_results(results, |iter| {
+                iter.take_while(|(key_bytes, _)| SlotColumnId::has_prefix(key_bytes))
+                    .map(|(k, v)| (k.into_owned(), v))
+                    .unzip()
+            })?;
+
+        for column_bytes in columns_to_remove {
+            let DataColumnIdentifier { block_root, index } =
+                DataColumnIdentifier::from_ssz_default(column_bytes)?;
+
+            keys_to_remove.push(
+                DataColumnSidecarByColumnId(block_root, index)
+                    .to_string()
+                    .into(),
+            )
         }
 
-        for column_id in columns_to_remove {
-            let DataColumnIdentifier { block_root, index } = column_id;
-            let key = DataColumnSidecarByColumnId(block_root, index).to_string();
-
-            self.database.delete(key)?;
-        }
-
-        for key in keys_to_remove {
-            self.database.delete(key)?;
-        }
-
-        Ok(())
+        self.database.delete_batch(keys_to_remove)
     }
 
     pub(crate) fn checkpoint_state_slot(&self) -> Result<Option<Slot>> {
@@ -874,15 +834,15 @@ impl<P: Preset> Storage<P> {
             .database
             .iterator_descending(..=BlockRootBySlot(start_from_slot).to_string())?;
 
-        let mut block_roots = vec![];
-
-        for result in results {
-            let (key_bytes, value_bytes) = result?;
+        let results = itertools::process_results(results, |iter| {
+            iter.take_while(|(key_bytes, _)| BlockRootBySlot::has_prefix(key_bytes))
+                .map(|(_, v)| v)
+                .collect::<Vec<_>>()
+        })?;
 
-            if !BlockRootBySlot::has_prefix(&key_bytes) {
-                break;
-            }
+        let mut block_roots = vec![];
 
+        for value_bytes in results {
             let block_root = H256::from_ssz_default(value_bytes)?;
 
             if self.contains_key(StateByBlockRoot(block_root))? {