Shrink persisted fork choice data (#7805)

Closes:

- #7760


  - [x] Remove `balances_cache` from `PersistedForkChoiceStore` (~65 MB saving on mainnet)
- [x] Remove `justified_balances` from `PersistedForkChoiceStore` (~16 MB saving on mainnet)
- [x] Remove `balances` from `ProtoArray`/`SszContainer`.
- [x] Implement zstd compression for votes
- [x] Fix bug in justified state usage
- [x] Bump schema version to V28 and implement migration.

Author: michaelsproul

Cargo.lock

@@ -1,4 +1,4 @@
- 
+
 [package]
 name = "beacon_chain"
 version = "0.2.0"
@@ -65,6 +65,7 @@ tracing = { workspace = true }
 tree_hash = { workspace = true }
 tree_hash_derive = { workspace = true }
 types = { workspace = true }
+zstd = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true }

beacon_node/beacon_chain/Cargo.toml

@@ -121,8 +121,8 @@ use std::sync::Arc;
 use std::time::Duration;
 use store::iter::{BlockRootsIterator, ParentRootBlockIterator, StateRootsIterator};
 use store::{
-    BlobSidecarListFromRoot, DatabaseBlock, Error as DBError, HotColdDB, HotStateSummary,
-    KeyValueStoreOp, StoreItem, StoreOp,
+    BlobSidecarListFromRoot, DBColumn, DatabaseBlock, Error as DBError, HotColdDB, HotStateSummary,
+    KeyValueStore, KeyValueStoreOp, StoreItem, StoreOp,
 };
 use task_executor::{ShutdownReason, TaskExecutor};
 use tokio_stream::Stream;
@@ -618,12 +618,15 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         reset_payload_statuses: ResetPayloadStatuses,
         spec: &ChainSpec,
     ) -> Result<Option<BeaconForkChoice<T>>, Error> {
-        let Some(persisted_fork_choice) =
-            store.get_item::<PersistedForkChoice>(&FORK_CHOICE_DB_KEY)?
+        let Some(persisted_fork_choice_bytes) = store
+            .hot_db
+            .get_bytes(DBColumn::ForkChoice, FORK_CHOICE_DB_KEY.as_slice())?
         else {
             return Ok(None);
         };
 
+        let persisted_fork_choice =
+            PersistedForkChoice::from_bytes(&persisted_fork_choice_bytes, store.get_config())?;
         let fc_store =
             BeaconForkChoiceStore::from_persisted(persisted_fork_choice.fork_choice_store, store)?;
 

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -27,6 +27,7 @@ pub enum Error {
     FailedToReadState(StoreError),
     MissingState(Hash256),
     BeaconStateError(BeaconStateError),
+    UnalignedCheckpoint { block_slot: Slot, state_slot: Slot },
     Arith(ArithError),
 }
 
@@ -136,7 +137,9 @@ pub struct BeaconForkChoiceStore<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<
     finalized_checkpoint: Checkpoint,
     justified_checkpoint: Checkpoint,
     justified_balances: JustifiedBalances,
+    justified_state_root: Hash256,
     unrealized_justified_checkpoint: Checkpoint,
+    unrealized_justified_state_root: Hash256,
     unrealized_finalized_checkpoint: Checkpoint,
     proposer_boost_root: Hash256,
     equivocating_indices: BTreeSet<u64>,
@@ -162,30 +165,48 @@ where
     /// It is assumed that `anchor` is already persisted in `store`.
     pub fn get_forkchoice_store(
         store: Arc<HotColdDB<E, Hot, Cold>>,
-        anchor: &BeaconSnapshot<E>,
+        anchor: BeaconSnapshot<E>,
     ) -> Result<Self, Error> {
-        let anchor_state = &anchor.beacon_state;
+        let unadvanced_state_root = anchor.beacon_state_root();
+        let mut anchor_state = anchor.beacon_state;
         let mut anchor_block_header = anchor_state.latest_block_header().clone();
-        if anchor_block_header.state_root == Hash256::zero() {
-            anchor_block_header.state_root = anchor.beacon_state_root();
+
+        // The anchor state MUST be on an epoch boundary (it should be advanced by the caller).
+        if !anchor_state
+            .slot()
+            .as_u64()
+            .is_multiple_of(E::slots_per_epoch())
+        {
+            return Err(Error::UnalignedCheckpoint {
+                block_slot: anchor_block_header.slot,
+                state_slot: anchor_state.slot(),
+            });
+        }
+
+        // Compute the accurate block root for the checkpoint block.
+        if anchor_block_header.state_root.is_zero() {
+            anchor_block_header.state_root = unadvanced_state_root;
         }
-        let anchor_root = anchor_block_header.canonical_root();
+        let anchor_block_root = anchor_block_header.canonical_root();
         let anchor_epoch = anchor_state.current_epoch();
         let justified_checkpoint = Checkpoint {
             epoch: anchor_epoch,
-            root: anchor_root,
+            root: anchor_block_root,
         };
         let finalized_checkpoint = justified_checkpoint;
-        let justified_balances = JustifiedBalances::from_justified_state(anchor_state)?;
+        let justified_balances = JustifiedBalances::from_justified_state(&anchor_state)?;
+        let justified_state_root = anchor_state.canonical_root()?;
 
         Ok(Self {
             store,
             balances_cache: <_>::default(),
             time: anchor_state.slot(),
             justified_checkpoint,
             justified_balances,
+            justified_state_root,
             finalized_checkpoint,
             unrealized_justified_checkpoint: justified_checkpoint,
+            unrealized_justified_state_root: justified_state_root,
             unrealized_finalized_checkpoint: finalized_checkpoint,
             proposer_boost_root: Hash256::zero(),
             equivocating_indices: BTreeSet::new(),
@@ -197,33 +218,72 @@ where
     /// on-disk database.
     pub fn to_persisted(&self) -> PersistedForkChoiceStore {
         PersistedForkChoiceStore {
-            balances_cache: self.balances_cache.clone(),
             time: self.time,
             finalized_checkpoint: self.finalized_checkpoint,
             justified_checkpoint: self.justified_checkpoint,
-            justified_balances: self.justified_balances.effective_balances.clone(),
+            justified_state_root: self.justified_state_root,
             unrealized_justified_checkpoint: self.unrealized_justified_checkpoint,
+            unrealized_justified_state_root: self.unrealized_justified_state_root,
             unrealized_finalized_checkpoint: self.unrealized_finalized_checkpoint,
             proposer_boost_root: self.proposer_boost_root,
             equivocating_indices: self.equivocating_indices.clone(),
         }
     }
 
     /// Restore `Self` from a previously-generated `PersistedForkChoiceStore`.
-    pub fn from_persisted(
-        persisted: PersistedForkChoiceStore,
+    ///
+    /// DEPRECATED. Can be deleted once migrations no longer require it.
+    pub fn from_persisted_v17(
+        persisted: PersistedForkChoiceStoreV17,
+        justified_state_root: Hash256,
+        unrealized_justified_state_root: Hash256,
         store: Arc<HotColdDB<E, Hot, Cold>>,
     ) -> Result<Self, Error> {
         let justified_balances =
             JustifiedBalances::from_effective_balances(persisted.justified_balances)?;
+
         Ok(Self {
             store,
-            balances_cache: persisted.balances_cache,
+            balances_cache: <_>::default(),
             time: persisted.time,
             finalized_checkpoint: persisted.finalized_checkpoint,
             justified_checkpoint: persisted.justified_checkpoint,
             justified_balances,
+            justified_state_root,
             unrealized_justified_checkpoint: persisted.unrealized_justified_checkpoint,
+            unrealized_justified_state_root,
+            unrealized_finalized_checkpoint: persisted.unrealized_finalized_checkpoint,
+            proposer_boost_root: persisted.proposer_boost_root,
+            equivocating_indices: persisted.equivocating_indices,
+            _phantom: PhantomData,
+        })
+    }
+
+    /// Restore `Self` from a previously-generated `PersistedForkChoiceStore`.
+    pub fn from_persisted(
+        persisted: PersistedForkChoiceStore,
+        store: Arc<HotColdDB<E, Hot, Cold>>,
+    ) -> Result<Self, Error> {
+        let justified_checkpoint = persisted.justified_checkpoint;
+        let justified_state_root = persisted.justified_state_root;
+
+        let update_cache = true;
+        let justified_state = store
+            .get_hot_state(&justified_state_root, update_cache)
+            .map_err(Error::FailedToReadState)?
+            .ok_or(Error::MissingState(justified_state_root))?;
+
+        let justified_balances = JustifiedBalances::from_justified_state(&justified_state)?;
+        Ok(Self {
+            store,
+            balances_cache: <_>::default(),
+            time: persisted.time,
+            finalized_checkpoint: persisted.finalized_checkpoint,
+            justified_checkpoint,
+            justified_balances,
+            justified_state_root,
+            unrealized_justified_checkpoint: persisted.unrealized_justified_checkpoint,
+            unrealized_justified_state_root: persisted.unrealized_justified_state_root,
             unrealized_finalized_checkpoint: persisted.unrealized_finalized_checkpoint,
             proposer_boost_root: persisted.proposer_boost_root,
             equivocating_indices: persisted.equivocating_indices,
@@ -261,6 +321,10 @@ where
         &self.justified_checkpoint
     }
 
+    fn justified_state_root(&self) -> Hash256 {
+        self.justified_state_root
+    }
+
     fn justified_balances(&self) -> &JustifiedBalances {
         &self.justified_balances
     }
@@ -273,6 +337,10 @@ where
         &self.unrealized_justified_checkpoint
     }
 
+    fn unrealized_justified_state_root(&self) -> Hash256 {
+        self.unrealized_justified_state_root
+    }
+
     fn unrealized_finalized_checkpoint(&self) -> &Checkpoint {
         &self.unrealized_finalized_checkpoint
     }
@@ -285,8 +353,13 @@ where
         self.finalized_checkpoint = checkpoint
     }
 
-    fn set_justified_checkpoint(&mut self, checkpoint: Checkpoint) -> Result<(), Error> {
+    fn set_justified_checkpoint(
+        &mut self,
+        checkpoint: Checkpoint,
+        justified_state_root: Hash256,
+    ) -> Result<(), Error> {
         self.justified_checkpoint = checkpoint;
+        self.justified_state_root = justified_state_root;
 
         if let Some(balances) = self.balances_cache.get(
             self.justified_checkpoint.root,
@@ -297,36 +370,24 @@ where
             self.justified_balances = JustifiedBalances::from_effective_balances(balances)?;
         } else {
             metrics::inc_counter(&metrics::BALANCES_CACHE_MISSES);
-            let justified_block = self
-                .store
-                .get_blinded_block(&self.justified_checkpoint.root)
-                .map_err(Error::FailedToReadBlock)?
-                .ok_or(Error::MissingBlock(self.justified_checkpoint.root))?
-                .deconstruct()
-                .0;
-
-            let max_slot = self
-                .justified_checkpoint
-                .epoch
-                .start_slot(E::slots_per_epoch());
-            let (_, state) = self
+
+            // Justified state is reasonably useful to cache, it might be finalized soon.
+            let update_cache = true;
+            let state = self
                 .store
-                .get_advanced_hot_state(
-                    self.justified_checkpoint.root,
-                    max_slot,
-                    justified_block.state_root(),
-                )
+                .get_hot_state(&self.justified_state_root, update_cache)
                 .map_err(Error::FailedToReadState)?
-                .ok_or_else(|| Error::MissingState(justified_block.state_root()))?;
+                .ok_or_else(|| Error::MissingState(self.justified_state_root))?;
 
             self.justified_balances = JustifiedBalances::from_justified_state(&state)?;
         }
 
         Ok(())
     }
 
-    fn set_unrealized_justified_checkpoint(&mut self, checkpoint: Checkpoint) {
+    fn set_unrealized_justified_checkpoint(&mut self, checkpoint: Checkpoint, state_root: Hash256) {
         self.unrealized_justified_checkpoint = checkpoint;
+        self.unrealized_justified_state_root = state_root;
     }
 
     fn set_unrealized_finalized_checkpoint(&mut self, checkpoint: Checkpoint) {
@@ -346,18 +407,48 @@ where
     }
 }
 
-pub type PersistedForkChoiceStore = PersistedForkChoiceStoreV17;
+pub type PersistedForkChoiceStore = PersistedForkChoiceStoreV28;
 
 /// A container which allows persisting the `BeaconForkChoiceStore` to the on-disk database.
-#[superstruct(variants(V17), variant_attributes(derive(Encode, Decode)), no_enum)]
+#[superstruct(
+    variants(V17, V28),
+    variant_attributes(derive(Encode, Decode)),
+    no_enum
+)]
 pub struct PersistedForkChoiceStore {
+    /// The balances cache was removed from disk storage in schema V28.
+    #[superstruct(only(V17))]
     pub balances_cache: BalancesCacheV8,
     pub time: Slot,
     pub finalized_checkpoint: Checkpoint,
     pub justified_checkpoint: Checkpoint,
+    /// The justified balances were removed from disk storage in schema V28.
+    #[superstruct(only(V17))]
     pub justified_balances: Vec<u64>,
+    /// The justified state root is stored so that it can be used to load the justified balances.
+    #[superstruct(only(V28))]
+    pub justified_state_root: Hash256,
     pub unrealized_justified_checkpoint: Checkpoint,
+    #[superstruct(only(V28))]
+    pub unrealized_justified_state_root: Hash256,
     pub unrealized_finalized_checkpoint: Checkpoint,
     pub proposer_boost_root: Hash256,
     pub equivocating_indices: BTreeSet<u64>,
 }
+
+// Convert V28 to V17 by adding balances and removing justified state roots.
+impl From<(PersistedForkChoiceStoreV28, JustifiedBalances)> for PersistedForkChoiceStoreV17 {
+    fn from((v28, balances): (PersistedForkChoiceStoreV28, JustifiedBalances)) -> Self {
+        Self {
+            balances_cache: Default::default(),
+            time: v28.time,
+            finalized_checkpoint: v28.finalized_checkpoint,
+            justified_checkpoint: v28.justified_checkpoint,
+            justified_balances: balances.effective_balances,
+            unrealized_justified_checkpoint: v28.unrealized_justified_checkpoint,
+            unrealized_finalized_checkpoint: v28.unrealized_finalized_checkpoint,
+            proposer_boost_root: v28.proposer_boost_root,
+            equivocating_indices: v28.equivocating_indices,
+        }
+    }
+}

beacon_node/beacon_chain/src/beacon_fork_choice_store.rs

@@ -394,7 +394,7 @@ where
                 .map_err(|e| format!("Failed to initialize genesis data column info: {:?}", e))?,
         );
 
-        let fc_store = BeaconForkChoiceStore::get_forkchoice_store(store, &genesis)
+        let fc_store = BeaconForkChoiceStore::get_forkchoice_store(store, genesis.clone())
             .map_err(|e| format!("Unable to initialize fork choice store: {e:?}"))?;
         let current_slot = None;
 
@@ -616,7 +616,7 @@ where
             beacon_state: weak_subj_state,
         };
 
-        let fc_store = BeaconForkChoiceStore::get_forkchoice_store(store, &snapshot)
+        let fc_store = BeaconForkChoiceStore::get_forkchoice_store(store, snapshot.clone())
             .map_err(|e| format!("Unable to initialize fork choice store: {e:?}"))?;
 
         let fork_choice = ForkChoice::from_anchor(
@@ -887,8 +887,9 @@ where
         self.pending_io_batch.push(BeaconChain::<
             Witness<TSlotClock,  E, THotStore, TColdStore>,
         >::persist_fork_choice_in_batch_standalone(
-            &fork_choice
-        ));
+            &fork_choice,
+            store.get_config(),
+        ).map_err(|e| format!("Fork choice compression error: {e:?}"))?);
         store
             .hot_db
             .do_atomically(self.pending_io_batch)