Persist fork choice after pruning

sigp · Oct 19, 2020 · 2c1b957 · 2c1b957
1 parent 7645634
commit 2c1b957
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -282,14 +282,19 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             .as_kv_store_op(BEACON_CHAIN_DB_KEY))
     }
 
-    /// Return a database operation for writing fork choice to disk.
-    pub fn persist_fork_choice_in_batch(&self) -> KeyValueStoreOp {
+    /// Return a `PersistedForkChoice` for writing fork choice to disk.
+    pub fn make_persisted_fork_choice(&self) -> PersistedForkChoice {
         let fork_choice = self.fork_choice.read();
-        let persisted_fork_choice = PersistedForkChoice {
+        PersistedForkChoice {
             fork_choice: fork_choice.to_persisted(),
             fork_choice_store: fork_choice.fc_store().to_persisted(),
-        };
-        persisted_fork_choice.as_kv_store_op(FORK_CHOICE_DB_KEY)
+        }
+    }
+
+    /// Return a database operation for writing fork choice to disk.
+    pub fn persist_fork_choice_in_batch(&self) -> KeyValueStoreOp {
+        self.make_persisted_fork_choice()
+            .as_kv_store_op(FORK_CHOICE_DB_KEY)
     }
 
     /// Persists `self.op_pool` to disk.
@@ -2118,6 +2123,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             new_finalized_checkpoint,
             self.head_tracker.clone(),
             self.make_persisted_head()?,
+            self.make_persisted_fork_choice(),
         )?;
 
         let _ = self.event_handler.register(EventKind::BeaconFinalization {

diff --git a/beacon_node/beacon_chain/src/migrate.rs b/beacon_node/beacon_chain/src/migrate.rs
@@ -1,7 +1,8 @@
-use crate::beacon_chain::BEACON_CHAIN_DB_KEY;
+use crate::beacon_chain::{BEACON_CHAIN_DB_KEY, FORK_CHOICE_DB_KEY};
 use crate::errors::BeaconChainError;
 use crate::head_tracker::HeadTracker;
 use crate::persisted_beacon_chain::PersistedBeaconChain;
+use crate::persisted_fork_choice::PersistedForkChoice;
 use parking_lot::Mutex;
 use slog::{debug, warn, Logger};
 use std::collections::{HashMap, HashSet};
@@ -74,6 +75,7 @@ pub struct MigrationNotification<E: EthSpec> {
     head_tracker: Arc<HeadTracker>,
     latest_checkpoint: Arc<Mutex<Checkpoint>>,
     persist_head: PersistedBeaconChain,
+    persist_fork_choice: PersistedForkChoice,
 }
 
 impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Hot, Cold> {
@@ -108,6 +110,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
         finalized_checkpoint: Checkpoint,
         head_tracker: Arc<HeadTracker>,
         persist_head: PersistedBeaconChain,
+        persist_fork_choice: PersistedForkChoice,
     ) -> Result<(), BeaconChainError> {
         let notif = MigrationNotification {
             finalized_state_root,
@@ -116,6 +119,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
             head_tracker,
             latest_checkpoint: self.latest_checkpoint.clone(),
             persist_head,
+            persist_fork_choice,
         };
 
         // Async path, on the background thread.
@@ -169,6 +173,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
             *latest_checkpoint,
             notif.finalized_checkpoint,
             notif.persist_head,
+            notif.persist_fork_choice,
             log,
         ) {
             Ok(PruningOutcome::DeferredConcurrentMutation) => {
@@ -240,6 +245,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
         old_finalized_checkpoint: Checkpoint,
         new_finalized_checkpoint: Checkpoint,
         mut persist_head: PersistedBeaconChain,
+        persist_fork_choice: PersistedForkChoice,
         log: &Logger,
     ) -> Result<PruningOutcome, BeaconChainError> {
         let old_finalized_slot = old_finalized_checkpoint
@@ -446,6 +452,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
         // but that's a price we're willing to pay (for now).
         persist_head.ssz_head_tracker = head_tracker.to_ssz_container();
         kv_batch.push(persist_head.as_kv_store_op(BEACON_CHAIN_DB_KEY));
+        kv_batch.push(persist_fork_choice.as_kv_store_op(FORK_CHOICE_DB_KEY));
 
         store.hot_db.do_atomically(kv_batch)?;
         debug!(log, "Database pruning complete");