Store state in correct column

dapplion · Dec 23, 2024 · 80eb882 · 80eb882
1 parent cfb8a3e
commit 80eb882
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 19 deletions.
diff --git a/beacon_node/store/src/hot_cold_store.rs b/beacon_node/store/src/hot_cold_store.rs
@@ -2,7 +2,6 @@ use crate::config::{OnDiskStoreConfig, StoreConfig};
 use crate::forwards_iter::{HybridForwardsBlockRootsIterator, HybridForwardsStateRootsIterator};
 use crate::hdiff::{HDiff, HDiffBuffer, HierarchyModuli, StorageStrategy};
 use crate::historic_state_cache::HistoricStateCache;
-use crate::impls::beacon_state::store_full_state;
 use crate::iter::{BlockRootsIterator, ParentRootBlockIterator, RootsIterator};
 use crate::leveldb_store::{BytesKey, LevelDB};
 use crate::memory_store::MemoryStore;
@@ -1527,7 +1526,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
                     "strategy" => "snapshot",
                     "slot" => slot,
                 );
-                store_full_state(state_root, state, ops)?;
+                self.store_hot_state_as_snapshot(state_root, state, ops)?;
             }
             StorageStrategy::DiffFrom(from_slot) => {
                 debug!(
@@ -1826,6 +1825,31 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
         }
     }
 
+    pub fn store_hot_state_as_snapshot(
+        &self,
+        state_root: &Hash256,
+        state: &BeaconState<E>,
+        ops: &mut Vec<KeyValueStoreOp>,
+    ) -> Result<(), Error> {
+        let bytes = state.as_ssz_bytes();
+        let compressed_value = {
+            let _timer = metrics::start_timer(&metrics::STORE_BEACON_STATE_FREEZER_COMPRESS_TIME);
+            let mut out = Vec::with_capacity(self.config.estimate_compressed_size(bytes.len()));
+            let mut encoder = Encoder::new(&mut out, self.config.compression_level)
+                .map_err(Error::Compression)?;
+            encoder.write_all(&bytes).map_err(Error::Compression)?;
+            encoder.finish().map_err(Error::Compression)?;
+            out
+        };
+
+        let key = get_key_for_col(
+            DBColumn::BeaconStateHotSnapshot.into(),
+            state_root.as_slice(),
+        );
+        ops.push(KeyValueStoreOp::PutKeyValue(key, compressed_value));
+        Ok(())
+    }
+
     fn load_hot_state_bytes_as_snapshot(
         &self,
         state_root: Hash256,

diff --git a/beacon_node/store/src/impls/beacon_state.rs b/beacon_node/store/src/impls/beacon_state.rs
@@ -1,23 +1,7 @@
 use crate::*;
-use ssz::{DecodeError, Encode};
+use ssz::DecodeError;
 use ssz_derive::Encode;
 
-pub fn store_full_state<E: EthSpec>(
-    state_root: &Hash256,
-    state: &BeaconState<E>,
-    ops: &mut Vec<KeyValueStoreOp>,
-) -> Result<(), Error> {
-    let bytes = {
-        let _overhead_timer = metrics::start_timer(&metrics::BEACON_STATE_WRITE_OVERHEAD_TIMES);
-        StorageContainer::new(state).as_ssz_bytes()
-    };
-    metrics::inc_counter_by(&metrics::BEACON_STATE_WRITE_BYTES, bytes.len() as u64);
-    metrics::inc_counter(&metrics::BEACON_STATE_WRITE_COUNT);
-    let key = get_key_for_col(DBColumn::BeaconState.into(), state_root.as_slice());
-    ops.push(KeyValueStoreOp::PutKeyValue(key, bytes));
-    Ok(())
-}
-
 // FIXME(tree-states): delete/move to migration
 pub fn get_full_state<KV: KeyValueStore<E>, E: EthSpec>(
     db: &KV,