fix: handling accessing shard db after shard removal (#2237)

* Address the issue that shard storage db is accessed after the db is removed

* Address unit test failure; cleanup code; enhance simtest

Author: halfprice

crates/typed-store/src/rocks.rs

@@ -586,22 +586,7 @@ impl<K, V> DBMap<K, V> {
         let from_buf = be_fix_int_ser(start)?;
         let to_buf = be_fix_int_ser(end)?;
         self.rocksdb
-            .compact_range_cf(&self.cf(), Some(from_buf), Some(to_buf));
-        Ok(())
-    }
-
-    /// Compact a range of keys in a specific column family
-    pub fn compact_range_raw(
-        &self,
-        cf_name: &str,
-        start: Vec<u8>,
-        end: Vec<u8>,
-    ) -> Result<(), TypedStoreError> {
-        let cf = self
-            .rocksdb
-            .cf_handle(cf_name)
-            .expect("compact range: column family does not exist");
-        self.rocksdb.compact_range_cf(&cf, Some(start), Some(end));
+            .compact_range_cf(&self.cf()?, Some(from_buf), Some(to_buf));
         Ok(())
     }
 
@@ -614,37 +599,24 @@ impl<K, V> DBMap<K, V> {
         let from_buf = be_fix_int_ser(start)?;
         let to_buf = be_fix_int_ser(end)?;
         self.rocksdb
-            .compact_range_to_bottom(&self.cf(), Some(from_buf), Some(to_buf));
+            .compact_range_to_bottom(&self.cf()?, Some(from_buf), Some(to_buf));
         Ok(())
     }
 
     /// Get the column family
-    pub fn cf(&self) -> Arc<rocksdb::BoundColumnFamily<'_>> {
-        self.rocksdb.cf_handle(&self.cf).unwrap_or_else(|| {
-            // Force capture backtrace regardless of RUST_BACKTRACE env var
-            let backtrace = std::backtrace::Backtrace::force_capture();
-            eprintln!("PANIC: Column family '{}' not found!", &self.cf);
-            eprintln!("Database path: {:?}", &self.rocksdb.db_path);
-            eprintln!("Full backtrace:\n{}", backtrace);
-            panic!(
-                "Map-keying column family {} should have been checked at DB creation",
-                &self.cf
-            )
-        })
+    pub fn cf(&self) -> Result<Arc<rocksdb::BoundColumnFamily<'_>>, TypedStoreError> {
+        self.rocksdb
+            .cf_handle(&self.cf)
+            .ok_or_else(|| TypedStoreError::UnregisteredColumn(self.cf.clone()))
     }
 
     /// Flush the column family
     pub fn flush(&self) -> Result<(), TypedStoreError> {
         self.rocksdb
-            .flush_cf(&self.cf())
+            .flush_cf(&self.cf()?)
             .map_err(|e| TypedStoreError::RocksDBError(e.into_string()))
     }
 
-    /// Set the options for the column family
-    pub fn set_options(&self, opts: &[(&str, &str)]) -> Result<(), rocksdb::Error> {
-        self.rocksdb.set_options_cf(&self.cf(), opts)
-    }
-
     fn get_int_property(
         rocksdb: &RocksDB,
         cf: &impl AsColumnFamilyRef,
@@ -686,7 +658,7 @@ impl<K, V> DBMap<K, V> {
         let results: Result<Vec<_>, TypedStoreError> = self
             .rocksdb
             .batched_multi_get_cf_opt(
-                &self.cf(),
+                &self.cf()?,
                 keys_refs,
                 /*sorted_keys=*/ false,
                 &self.opts.readopts(),
@@ -940,7 +912,7 @@ impl<K, V> DBMap<K, V> {
             .expect("the function parameters are valid");
         let mut value_hist = hdrhistogram::Histogram::<u64>::new_with_max(100000, 2)
             .expect("the function parameters are valid");
-        for item in self.safe_iter() {
+        for item in self.safe_iter()? {
             let (key, value) = item?;
             num_keys += 1;
             let key_len = be_fix_int_ser(key.borrow())?.len();
@@ -1043,7 +1015,7 @@ impl<K, V> DBMap<K, V> {
                 .unwrap_or(Bound::Unbounded),
         ));
 
-        let db_iter = self.rocksdb.raw_iterator_cf(&self.cf(), readopts);
+        let db_iter = self.rocksdb.raw_iterator_cf(&self.cf()?, readopts);
         let (_timer, bytes_scanned, keys_scanned, _perf_ctx) = self.create_iter_context();
         let iter = SafeIter::new(
             self.cf.clone(),
@@ -1270,7 +1242,7 @@ impl DBBatch {
             .into_iter()
             .try_for_each::<_, Result<_, TypedStoreError>>(|k| {
                 let k_buf = be_fix_int_ser(k.borrow())?;
-                self.batch.delete_cf(&db.cf(), k_buf);
+                self.batch.delete_cf(&db.cf()?, k_buf);
 
                 Ok(())
             })?;
@@ -1299,7 +1271,7 @@ impl DBBatch {
         let from_buf = be_fix_int_ser(from)?;
         let to_buf = be_fix_int_ser(to)?;
 
-        self.batch.delete_range_cf(&db.cf(), from_buf, to_buf);
+        self.batch.delete_range_cf(&db.cf()?, from_buf, to_buf);
         Ok(())
     }
 
@@ -1319,7 +1291,7 @@ impl DBBatch {
                 let k_buf = be_fix_int_ser(k.borrow())?;
                 let v_buf = bcs::to_bytes(v.borrow()).map_err(typed_store_err_from_bcs_err)?;
                 total += k_buf.len() + v_buf.len();
-                self.batch.put_cf(&db.cf(), k_buf, v_buf);
+                self.batch.put_cf(&db.cf()?, k_buf, v_buf);
                 Ok(())
             })?;
         self.db_metrics
@@ -1343,7 +1315,7 @@ impl DBBatch {
             .into_iter()
             .try_for_each::<_, Result<_, TypedStoreError>>(|(k, v)| {
                 let k_buf = be_fix_int_ser(k.borrow())?;
-                self.batch.merge_cf(&db.cf(), k_buf, v);
+                self.batch.merge_cf(&db.cf()?, k_buf, v);
                 Ok(())
             })?;
         Ok(self)
@@ -1370,10 +1342,10 @@ where
         let readopts = self.opts.readopts();
         Ok(self
             .rocksdb
-            .key_may_exist_cf(&self.cf(), &key_buf, &readopts)
+            .key_may_exist_cf(&self.cf()?, &key_buf, &readopts)
             && self
                 .rocksdb
-                .get_pinned_cf_opt(&self.cf(), &key_buf, &readopts)
+                .get_pinned_cf_opt(&self.cf()?, &key_buf, &readopts)
                 .map_err(typed_store_err_from_rocks_err)?
                 .is_some())
     }
@@ -1406,7 +1378,7 @@ where
         let key_buf = be_fix_int_ser(key)?;
         let res = self
             .rocksdb
-            .get_pinned_cf_opt(&self.cf(), &key_buf, &self.opts.readopts())
+            .get_pinned_cf_opt(&self.cf()?, &key_buf, &self.opts.readopts())
             .map_err(typed_store_err_from_rocks_err)?;
         self.db_metrics
             .op_metrics
@@ -1452,7 +1424,7 @@ where
                 .report_metrics(&self.cf);
         }
         self.rocksdb
-            .put_cf(&self.cf(), &key_buf, &value_buf, &self.opts.writeopts())
+            .put_cf(&self.cf()?, &key_buf, &value_buf, &self.opts.writeopts())
             .map_err(typed_store_err_from_rocks_err)?;
 
         let elapsed = timer.stop_and_record();
@@ -1488,7 +1460,7 @@ where
         };
         let key_buf = be_fix_int_ser(key)?;
         self.rocksdb
-            .delete_cf(&self.cf(), key_buf, &self.opts.writeopts())
+            .delete_cf(&self.cf()?, key_buf, &self.opts.writeopts())
             .map_err(typed_store_err_from_rocks_err)?;
         self.db_metrics
             .op_metrics
@@ -1526,7 +1498,7 @@ where
     /// overridden in the config), so please use this function with caution
     #[tracing::instrument(level = "trace", skip_all, err)]
     fn schedule_delete_all(&self) -> Result<(), TypedStoreError> {
-        let first_key = self.safe_iter().next().transpose()?.map(|(k, _v)| k);
+        let first_key = self.safe_iter()?.next().transpose()?.map(|(k, _v)| k);
         let last_key = self
             .reversed_safe_iter_with_bounds(None, None)?
             .next()
@@ -1541,57 +1513,63 @@ where
     }
 
     fn is_empty(&self) -> bool {
-        self.safe_iter().next().is_none()
+        self.safe_iter()
+            .expect("safe_iter should not fail")
+            .next()
+            .is_none()
     }
 
-    fn safe_iter(&'a self) -> Self::SafeIterator {
+    fn safe_iter(&'a self) -> Result<Self::SafeIterator, TypedStoreError> {
         let db_iter = self
             .rocksdb
-            .raw_iterator_cf(&self.cf(), self.opts.readopts());
+            .raw_iterator_cf(&self.cf()?, self.opts.readopts());
         let (_timer, bytes_scanned, keys_scanned, _perf_ctx) = self.create_iter_context();
-        SafeIter::new(
+        Ok(SafeIter::new(
             self.cf.clone(),
             db_iter,
             _timer,
             _perf_ctx,
             bytes_scanned,
             keys_scanned,
             Some(self.db_metrics.clone()),
-        )
+        ))
     }
 
     fn safe_iter_with_bounds(
         &'a self,
         lower_bound: Option<K>,
         upper_bound: Option<K>,
-    ) -> Self::SafeIterator {
+    ) -> Result<Self::SafeIterator, TypedStoreError> {
         let readopts = self.create_read_options_with_bounds(lower_bound, upper_bound);
-        let db_iter = self.rocksdb.raw_iterator_cf(&self.cf(), readopts);
+        let db_iter = self.rocksdb.raw_iterator_cf(&self.cf()?, readopts);
         let (_timer, bytes_scanned, keys_scanned, _perf_ctx) = self.create_iter_context();
-        SafeIter::new(
+        Ok(SafeIter::new(
             self.cf.clone(),
             db_iter,
             _timer,
             _perf_ctx,
             bytes_scanned,
             keys_scanned,
             Some(self.db_metrics.clone()),
-        )
+        ))
     }
 
-    fn safe_range_iter(&'a self, range: impl RangeBounds<K>) -> Self::SafeIterator {
+    fn safe_range_iter(
+        &'a self,
+        range: impl RangeBounds<K>,
+    ) -> Result<Self::SafeIterator, TypedStoreError> {
         let readopts = self.create_read_options_with_range(range);
-        let db_iter = self.rocksdb.raw_iterator_cf(&self.cf(), readopts);
+        let db_iter = self.rocksdb.raw_iterator_cf(&self.cf()?, readopts);
         let (_timer, bytes_scanned, keys_scanned, _perf_ctx) = self.create_iter_context();
-        SafeIter::new(
+        Ok(SafeIter::new(
             self.cf.clone(),
             db_iter,
             _timer,
             _perf_ctx,
             bytes_scanned,
             keys_scanned,
             Some(self.db_metrics.clone()),
-        )
+        ))
     }
 
     /// Returns a vector of values corresponding to the keys provided.

crates/typed-store/src/rocks/tests.rs

@@ -72,7 +72,7 @@ where
     K: Serialize + DeserializeOwned,
     V: Serialize + DeserializeOwned,
 {
-    TestIteratorWrapper::SafeIter(db.safe_iter())
+    TestIteratorWrapper::SafeIter(db.safe_iter().expect("failed to get iterator"))
 }
 
 fn get_reverse_iter<K, V>(
@@ -97,7 +97,10 @@ where
     K: Serialize + DeserializeOwned,
     V: Serialize + DeserializeOwned,
 {
-    TestIteratorWrapper::SafeIter(db.safe_iter_with_bounds(lower_bound, upper_bound))
+    TestIteratorWrapper::SafeIter(
+        db.safe_iter_with_bounds(lower_bound, upper_bound)
+            .expect("failed to get iterator"),
+    )
 }
 
 fn get_range_iter<K, V>(
@@ -108,7 +111,7 @@ where
     K: Serialize + DeserializeOwned,
     V: Serialize + DeserializeOwned,
 {
-    TestIteratorWrapper::SafeIter(db.safe_range_iter(range))
+    TestIteratorWrapper::SafeIter(db.safe_range_iter(range).expect("failed to get iterator"))
 }
 
 #[tokio::test]
@@ -489,17 +492,17 @@ async fn test_clear() {
     insert_batch.write().expect("Failed to execute batch");
 
     // Check we have multiple entries
-    assert!(db.safe_iter().count() > 1);
+    assert!(db.safe_iter().expect("failed to get iterator").count() > 1);
     let _ = db.unsafe_clear();
-    assert_eq!(db.safe_iter().count(), 0);
+    assert_eq!(db.safe_iter().expect("failed to get iterator").count(), 0);
     // Clear again to ensure safety when clearing empty map
     let _ = db.unsafe_clear();
-    assert_eq!(db.safe_iter().count(), 0);
+    assert_eq!(db.safe_iter().expect("failed to get iterator").count(), 0);
     // Clear with one item
     let _ = db.insert(&1, &"e".to_string());
-    assert_eq!(db.safe_iter().count(), 1);
+    assert_eq!(db.safe_iter().expect("failed to get iterator").count(), 1);
     let _ = db.unsafe_clear();
-    assert_eq!(db.safe_iter().count(), 0);
+    assert_eq!(db.safe_iter().expect("failed to get iterator").count(), 0);
 }
 
 #[tokio::test]
@@ -554,7 +557,9 @@ async fn test_iter_with_bounds() {
     );
 
     // Specify a bound outside of dataset.
-    let db_iter = db.safe_iter_with_bounds(Some(200), Some(300));
+    let db_iter = db
+        .safe_iter_with_bounds(Some(200), Some(300))
+        .expect("failed to get iterator");
     assert!(db_iter.collect::<Vec<_>>().is_empty());
 
     // Skip to first key in the bound (bound is [1, 50))
@@ -631,12 +636,12 @@ async fn test_is_empty() {
     insert_batch.write().expect("Failed to execute batch");
 
     // Check we have multiple entries and not empty
-    assert!(db.safe_iter().count() > 1);
+    assert!(db.safe_iter().expect("failed to get iterator").count() > 1);
     assert!(!db.is_empty());
 
     // Clear again to ensure empty works after clearing
     let _ = db.unsafe_clear();
-    assert_eq!(db.safe_iter().count(), 0);
+    assert_eq!(db.safe_iter().expect("failed to get iterator").count(), 0);
     assert!(db.is_empty());
 }
 
@@ -713,7 +718,10 @@ async fn test_multi_remove() {
     // Remove 50 items
     db.multi_remove(keys_vals.clone().map(|kv| kv.0).take(50))
         .expect("Failed to multi-remove");
-    assert_eq!(db.safe_iter().count(), 101 - 50);
+    assert_eq!(
+        db.safe_iter().expect("failed to get iterator").count(),
+        101 - 50
+    );
 
     // Check that the remaining are present
     for (k, v) in keys_vals.skip(50) {

crates/typed-store/src/traits.rs

@@ -53,17 +53,20 @@ where
     fn is_empty(&self) -> bool;
 
     /// Same as `iter` but performs status check.
-    fn safe_iter(&'a self) -> Self::SafeIterator;
+    fn safe_iter(&'a self) -> Result<Self::SafeIterator, Self::Error>;
 
     /// Same as `iter_with_bounds` but performs status check.
     fn safe_iter_with_bounds(
         &'a self,
         lower_bound: Option<K>,
         upper_bound: Option<K>,
-    ) -> Self::SafeIterator;
+    ) -> Result<Self::SafeIterator, Self::Error>;
 
     /// Same as `range_iter` but performs status check.
-    fn safe_range_iter(&'a self, range: impl RangeBounds<K>) -> Self::SafeIterator;
+    fn safe_range_iter(
+        &'a self,
+        range: impl RangeBounds<K>,
+    ) -> Result<Self::SafeIterator, Self::Error>;
 
     /// Returns a vector of values corresponding to the keys provided, non-atomically.
     fn multi_get<J>(&self, keys: impl IntoIterator<Item = J>) -> Result<Vec<Option<V>>, Self::Error>

crates/walrus-service/src/event/event_processor/processor.rs

@@ -210,7 +210,7 @@ impl EventProcessor {
     pub fn poll(&self, from: u64) -> Result<Vec<PositionedStreamEvent>, TypedStoreError> {
         self.stores
             .event_store
-            .safe_iter_with_bounds(Some(from), None)
+            .safe_iter_with_bounds(Some(from), None)?
             .take(MAX_EVENTS_PER_POLL)
             .map(|result| result.map(|(_, event)| event))
             .collect()
@@ -222,7 +222,7 @@ impl EventProcessor {
         let mut iter = self
             .stores
             .event_store
-            .safe_iter_with_bounds(Some(from), None);
+            .safe_iter_with_bounds(Some(from), None)?;
         let Some(result) = iter.next() else {
             return Ok(None);
         };
@@ -330,7 +330,7 @@ impl EventProcessor {
                         .with_fixint_encoding()
                         .serialize(&commit_index)?;
                     self.stores.event_store.rocksdb.delete_file_in_range(
-                        &self.stores.event_store.cf(),
+                        &self.stores.event_store.cf()?,
                         &start,
                         &end,
                     )?;