Unify fork choice with chain

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -19,8 +19,7 @@ use crate::observed_attestations::{Error as AttestationObservationError, Observe
 use crate::observed_attesters::{ObservedAggregators, ObservedAttesters};
 use crate::observed_block_producers::ObservedBlockProducers;
 use crate::observed_operations::{ObservationOutcome, ObservedOperations};
-use crate::persisted_beacon_chain::PersistedBeaconChain;
-use crate::persisted_fork_choice::PersistedForkChoice;
+use crate::persisted_beacon_chain::{PersistedBeaconChain, PersistedForkChoice};
 use crate::shuffling_cache::ShufflingCache;
 use crate::snapshot_cache::SnapshotCache;
 use crate::timeout_rw_lock::TimeoutRwLock;
@@ -242,33 +241,21 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     /// We want to ensure that the head never out dates the fork choice to avoid having references
     /// to blocks that do not exist in fork choice.
     pub fn persist_head_and_fork_choice(&self) -> Result<(), Error> {
-        let canonical_head = self
-            .canonical_head
-            .try_read_for(HEAD_LOCK_TIMEOUT)
-            .ok_or_else(|| Error::CanonicalHeadLockTimeout)?;
         let fork_choice = self.fork_choice.read();
 
-        let persisted_head = PersistedBeaconChain {
-            canonical_head_block_root: canonical_head.beacon_block_root,
-            genesis_block_root: self.genesis_block_root,
-            ssz_head_tracker: self.head_tracker.to_ssz_container(),
-        };
-
         let persisted_fork_choice = PersistedForkChoice {
             fork_choice: fork_choice.to_persisted(),
             fork_choice_store: fork_choice.fc_store().to_persisted(),
         };
 
-        drop(canonical_head);
-        drop(fork_choice);
+        let persisted_head = PersistedBeaconChain {
+            genesis_time: self.slot_clock.genesis_duration().as_secs(),
+            genesis_block_root: self.genesis_block_root,
+            ssz_head_tracker: self.head_tracker.to_ssz_container(),
+            persisted_fork_choice,
+        };
 
-        {
-            let _timer = metrics::start_timer(&metrics::PERSIST_FORK_CHOICE);
-            self.store.put_item(
-                &Hash256::from_slice(&FORK_CHOICE_DB_KEY),
-                &persisted_fork_choice,
-            )?;
-        }
+        drop(fork_choice);
 
         {
             let _timer = metrics::start_timer(&metrics::PERSIST_HEAD);

beacon_node/beacon_chain/src/builder.rs

@@ -1,12 +1,9 @@
-use crate::beacon_chain::{
-    BEACON_CHAIN_DB_KEY, ETH1_CACHE_DB_KEY, FORK_CHOICE_DB_KEY, OP_POOL_DB_KEY,
-};
+use crate::beacon_chain::{BEACON_CHAIN_DB_KEY, ETH1_CACHE_DB_KEY, OP_POOL_DB_KEY};
 use crate::eth1_chain::{CachingEth1Backend, SszEth1};
 use crate::events::NullEventHandler;
 use crate::head_tracker::HeadTracker;
 use crate::migrate::Migrate;
 use crate::persisted_beacon_chain::PersistedBeaconChain;
-use crate::persisted_fork_choice::PersistedForkChoice;
 use crate::shuffling_cache::ShufflingCache;
 use crate::snapshot_cache::{SnapshotCache, DEFAULT_SNAPSHOT_CACHE_SIZE};
 use crate::timeout_rw_lock::TimeoutRwLock;
@@ -294,11 +291,32 @@ where
                 .map_err(|e| format!("Failed to decode head tracker for database: {:?}", e))?,
         );
 
+        /*
+         * Fork choice
+         */
+
+        let fc_store = BeaconForkChoiceStore::from_persisted(
+            chain.persisted_fork_choice.fork_choice_store,
+            store.clone(),
+        )
+        .map_err(|e| format!("Unable to load ForkChoiceStore: {:?}", e))?;
+
+        let mut fork_choice =
+            ForkChoice::from_persisted(chain.persisted_fork_choice.fork_choice, fc_store)
+                .map_err(|e| format!("Unable to parse persisted fork choice from disk: {:?}", e))?;
+
+        let current_slot = chain.genesis_time / (self.spec.milliseconds_per_slot / 1_000);
+
+        let head_block_root = fork_choice
+            .get_head(current_slot.into())
+            .map_err(|e| format!("Unable to find head block: {:?}", e))?;
+
+        self.fork_choice = Some(fork_choice);
+
         /*
          * Canonical head
          */
 
-        let head_block_root = chain.canonical_head_block_root;
         let head_block = store
             .get_item::<SignedBeaconBlock<TEthSpec>>(&head_block_root)
             .map_err(|e| format!("DB error when reading head block: {:?}", e))?
@@ -337,26 +355,6 @@ where
 
         self.validator_pubkey_cache = Some(pubkey_cache);
 
-        /*
-         * Fork choice
-         */
-
-        let persisted_fork_choice = store
-            .get_item::<PersistedForkChoice>(&Hash256::from_slice(&FORK_CHOICE_DB_KEY))
-            .map_err(|e| format!("DB error when reading persisted fork choice: {:?}", e))?
-            .ok_or_else(|| "Fork choice not found in store".to_string())?;
-
-        let fc_store = BeaconForkChoiceStore::from_persisted(
-            persisted_fork_choice.fork_choice_store,
-            store.clone(),
-        )
-        .map_err(|e| format!("Unable to load ForkChoiceStore: {:?}", e))?;
-
-        self.fork_choice = Some(
-            ForkChoice::from_persisted(persisted_fork_choice.fork_choice, fc_store)
-                .map_err(|e| format!("Unable to parse persisted fork choice from disk: {:?}", e))?,
-        );
-
         Ok(self)
     }
 

beacon_node/beacon_chain/src/metrics.rs

@@ -244,8 +244,6 @@ lazy_static! {
         try_create_histogram("beacon_persist_op_pool", "Time taken to persist the operations pool");
     pub static ref PERSIST_ETH1_CACHE: Result<Histogram> =
         try_create_histogram("beacon_persist_eth1_cache", "Time taken to persist the eth1 caches");
-    pub static ref PERSIST_FORK_CHOICE: Result<Histogram> =
-        try_create_histogram("beacon_persist_fork_choice", "Time taken to persist the fork choice struct");
 
     /*
      * Eth1

beacon_node/beacon_chain/src/persisted_beacon_chain.rs

@@ -1,14 +1,17 @@
+use crate::beacon_fork_choice_store::PersistedForkChoiceStore as ForkChoiceStore;
 use crate::head_tracker::SszHeadTracker;
+use fork_choice::PersistedForkChoice as ForkChoice;
 use ssz::{Decode, Encode};
 use ssz_derive::{Decode, Encode};
 use store::{DBColumn, Error as StoreError, StoreItem};
 use types::Hash256;
 
-#[derive(Clone, Encode, Decode)]
+#[derive(Encode, Decode)]
 pub struct PersistedBeaconChain {
-    pub canonical_head_block_root: Hash256,
+    pub genesis_time: u64,
     pub genesis_block_root: Hash256,
     pub ssz_head_tracker: SszHeadTracker,
+    pub persisted_fork_choice: PersistedForkChoice,
 }
 
 impl StoreItem for PersistedBeaconChain {
@@ -24,3 +27,23 @@ impl StoreItem for PersistedBeaconChain {
         Self::from_ssz_bytes(bytes).map_err(Into::into)
     }
 }
+
+#[derive(Encode, Decode)]
+pub struct PersistedForkChoice {
+    pub fork_choice: ForkChoice,
+    pub fork_choice_store: ForkChoiceStore,
+}
+
+impl StoreItem for PersistedForkChoice {
+    fn db_column() -> DBColumn {
+        DBColumn::ForkChoice
+    }
+
+    fn as_store_bytes(&self) -> Vec<u8> {
+        self.as_ssz_bytes()
+    }
+
+    fn from_store_bytes(bytes: &[u8]) -> std::result::Result<Self, StoreError> {
+        Self::from_ssz_bytes(bytes).map_err(Into::into)
+    }
+}

common/slot_clock/src/lib.rs

@@ -52,6 +52,9 @@ pub trait SlotClock: Send + Sync + Sized {
     /// Returns the first slot to be returned at the genesis time.
     fn genesis_slot(&self) -> Slot;
 
+    /// Returns the genesis time, as a duration since UNIX epoch.
+    fn genesis_duration(&self) -> Duration;
+
     /// Returns the slot if the internal clock were advanced by `duration`.
     fn now_with_future_tolerance(&self, tolerance: Duration) -> Option<Slot> {
         self.slot_of(self.now_duration()?.checked_add(tolerance)?)

common/slot_clock/src/manual_slot_clock.rs

@@ -41,10 +41,6 @@ impl ManualSlotClock {
         self.set_slot(self.now().unwrap().as_u64() + 1)
     }
 
-    pub fn genesis_duration(&self) -> &Duration {
-        &self.genesis_duration
-    }
-
     /// Returns the duration between UNIX epoch and the start of `slot`.
     pub fn start_of(&self, slot: Slot) -> Option<Duration> {
         let slot = slot
@@ -143,6 +139,10 @@ impl SlotClock for ManualSlotClock {
         self.slot_duration
     }
 
+    fn genesis_duration(&self) -> Duration {
+        self.genesis_duration
+    }
+
     fn duration_to_slot(&self, slot: Slot) -> Option<Duration> {
         self.duration_to_slot(slot, *self.current_time.read())
     }

common/slot_clock/src/system_time_slot_clock.rs

@@ -24,7 +24,7 @@ impl SlotClock for SystemTimeSlotClock {
 
     fn is_prior_to_genesis(&self) -> Option<bool> {
         let now = SystemTime::now().duration_since(UNIX_EPOCH).ok()?;
-        Some(now < *self.clock.genesis_duration())
+        Some(now < self.clock.genesis_duration())
     }
 
     fn now_duration(&self) -> Option<Duration> {
@@ -49,6 +49,10 @@ impl SlotClock for SystemTimeSlotClock {
         self.clock.slot_duration()
     }
 
+    fn genesis_duration(&self) -> Duration {
+        self.clock.genesis_duration()
+    }
+
     fn duration_to_slot(&self, slot: Slot) -> Option<Duration> {
         let now = SystemTime::now().duration_since(UNIX_EPOCH).ok()?;
         self.clock.duration_to_slot(slot, now)