From 6973184b06017f19894ab0925898a205a2cfdace Mon Sep 17 00:00:00 2001
From: Eitan Seri-Levi <eserilev@ucsc.edu>
Date: Wed, 29 Jan 2025 12:22:21 +0300
Subject: [PATCH] Fix Redb implementation and add CI checks (#6856)

---
 Makefile                                    |  2 +-
 beacon_node/store/src/database/redb_impl.rs | 33 ++++++++++++---------
 beacon_node/store/src/forwards_iter.rs      |  1 +
 beacon_node/store/src/hot_cold_store.rs     |  6 ++--
 4 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/Makefile b/Makefile
index e8b44cb780..0f08afd168 100644
--- a/Makefile
+++ b/Makefile
@@ -222,7 +222,7 @@ lint-fix:
 
 # Also run the lints on the optimized-only tests
 lint-full:
-	RUSTFLAGS="-C debug-assertions=no $(RUSTFLAGS)" $(MAKE) lint
+	TEST_FEATURES="beacon-node-leveldb,beacon-node-redb,${TEST_FEATURES}"  RUSTFLAGS="-C debug-assertions=no $(RUSTFLAGS)" $(MAKE) lint
 
 # Runs the makefile in the `ef_tests` repo.
 #
diff --git a/beacon_node/store/src/database/redb_impl.rs b/beacon_node/store/src/database/redb_impl.rs
index 6a776da7b1..cbe575d184 100644
--- a/beacon_node/store/src/database/redb_impl.rs
+++ b/beacon_node/store/src/database/redb_impl.rs
@@ -215,11 +215,12 @@ impl<E: EthSpec> Redb<E> {
         let table_definition: TableDefinition<'_, &[u8], &[u8]> =
             TableDefinition::new(column.into());
 
-        let iter = {
+        let result = (|| {
             let open_db = self.db.read();
             let read_txn = open_db.begin_read()?;
             let table = read_txn.open_table(table_definition)?;
-            table.range(from..)?.map(move |res| {
+            let range = table.range(from..)?;
+            Ok(range.map(move |res| {
                 let (key, _) = res?;
                 metrics::inc_counter_vec(&metrics::DISK_DB_KEY_READ_COUNT, &[column.into()]);
                 metrics::inc_counter_vec_by(
@@ -228,10 +229,13 @@ impl<E: EthSpec> Redb<E> {
                     key.value().len() as u64,
                 );
                 K::from_bytes(key.value())
-            })
-        };
+            }))
+        })();
 
-        Box::new(iter)
+        match result {
+            Ok(iter) => Box::new(iter),
+            Err(err) => Box::new(std::iter::once(Err(err))),
+        }
     }
 
     /// Iterate through all keys and values in a particular column.
@@ -243,15 +247,13 @@ impl<E: EthSpec> Redb<E> {
         let table_definition: TableDefinition<'_, &[u8], &[u8]> =
             TableDefinition::new(column.into());
 
-        let prefix = from.to_vec();
-
-        let iter = {
+        let result = (|| {
             let open_db = self.db.read();
             let read_txn = open_db.begin_read()?;
             let table = read_txn.open_table(table_definition)?;
+            let range = table.range(from..)?;
 
-            table
-                .range(from..)?
+            Ok(range
                 .take_while(move |res| match res.as_ref() {
                     Ok((_, _)) => true,
                     Err(_) => false,
@@ -265,14 +267,17 @@ impl<E: EthSpec> Redb<E> {
                         value.value().len() as u64,
                     );
                     Ok((K::from_bytes(key.value())?, value.value().to_vec()))
-                })
-        };
+                }))
+        })();
 
-        Ok(Box::new(iter))
+        match result {
+            Ok(iter) => Box::new(iter),
+            Err(err) => Box::new(std::iter::once(Err(err))),
+        }
     }
 
     pub fn iter_column<K: Key>(&self, column: DBColumn) -> ColumnIter<K> {
-        self.iter_column_from(column, &vec![0; column.key_size()], |_, _| true)
+        self.iter_column_from(column, &vec![0; column.key_size()])
     }
 
     pub fn delete_batch(&self, col: DBColumn, ops: HashSet<&[u8]>) -> Result<(), Error> {
diff --git a/beacon_node/store/src/forwards_iter.rs b/beacon_node/store/src/forwards_iter.rs
index 5300a74c06..255b7d8eac 100644
--- a/beacon_node/store/src/forwards_iter.rs
+++ b/beacon_node/store/src/forwards_iter.rs
@@ -158,6 +158,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> Iterator
             return None;
         }
         self.inner
+            .as_mut()
             .next()?
             .and_then(|(slot_bytes, root_bytes)| {
                 let slot = slot_bytes
diff --git a/beacon_node/store/src/hot_cold_store.rs b/beacon_node/store/src/hot_cold_store.rs
index 75251cb5fb..45b1983492 100644
--- a/beacon_node/store/src/hot_cold_store.rs
+++ b/beacon_node/store/src/hot_cold_store.rs
@@ -14,8 +14,8 @@ use crate::metadata::{
 };
 use crate::state_cache::{PutStateOutcome, StateCache};
 use crate::{
-    get_data_column_key, metrics, parse_data_column_key, BlobSidecarListFromRoot, DBColumn,
-    DatabaseBlock, Error, ItemStore, KeyValueStore, KeyValueStoreOp, StoreItem, StoreOp,
+    get_data_column_key, metrics, parse_data_column_key, BlobSidecarListFromRoot, ColumnKeyIter,
+    DBColumn, DatabaseBlock, Error, ItemStore, KeyValueStore, KeyValueStoreOp, StoreItem, StoreOp,
 };
 use itertools::{process_results, Itertools};
 use lru::LruCache;
@@ -405,7 +405,7 @@ impl<E: EthSpec> HotColdDB<E, BeaconNodeBackend<E>, BeaconNodeBackend<E>> {
     }
 
     /// Return an iterator over the state roots of all temporary states.
-    pub fn iter_temporary_state_roots(&self) -> impl Iterator<Item = Result<Hash256, Error>> + '_ {
+    pub fn iter_temporary_state_roots(&self) -> ColumnKeyIter<Hash256> {
         self.hot_db
             .iter_column_keys::<Hash256>(DBColumn::BeaconStateTemporary)
     }