Add HeadTracker struct

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -3,6 +3,7 @@ use crate::errors::{BeaconChainError as Error, BlockProductionError};
 use crate::eth1_chain::{Eth1Chain, Eth1ChainBackend};
 use crate::events::{EventHandler, EventKind};
 use crate::fork_choice::{Error as ForkChoiceError, ForkChoice};
+use crate::head_tracker::HeadTracker;
 use crate::metrics;
 use crate::persisted_beacon_chain::{PersistedBeaconChain, BEACON_CHAIN_DB_KEY};
 use lmd_ghost::LmdGhost;
@@ -126,6 +127,8 @@ pub struct BeaconChain<T: BeaconChainTypes> {
     pub fork_choice: ForkChoice<T>,
     /// A handler for events generated by the beacon chain.
     pub event_handler: T::EventHandler,
+    /// Used to track the heads of the beacon chain.
+    pub(crate) head_tracker: HeadTracker,
     /// Logging to CLI, etc.
     pub(crate) log: Logger,
 }
@@ -141,6 +144,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             canonical_head: self.canonical_head.read().clone(),
             op_pool: PersistedOperationPool::from_operation_pool(&self.op_pool),
             genesis_block_root: self.genesis_block_root,
+            ssz_head_tracker: self.head_tracker.to_ssz_container(),
         };
 
         let key = Hash256::from_slice(&BEACON_CHAIN_DB_KEY.as_bytes());
@@ -1219,6 +1223,8 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
         metrics::stop_timer(db_write_timer);
 
+        self.head_tracker.register_block(block_root, &block);
+
         let fork_choice_register_timer =
             metrics::start_timer(&metrics::BLOCK_PROCESSING_FORK_CHOICE_REGISTER);
 

beacon_node/beacon_chain/src/builder.rs

@@ -1,5 +1,6 @@
 use crate::eth1_chain::CachingEth1Backend;
 use crate::events::NullEventHandler;
+use crate::head_tracker::HeadTracker;
 use crate::persisted_beacon_chain::{PersistedBeaconChain, BEACON_CHAIN_DB_KEY};
 use crate::{
     BeaconChain, BeaconChainTypes, CheckPoint, Eth1Chain, Eth1ChainBackend, EventHandler,
@@ -9,7 +10,7 @@ use eth1::Config as Eth1Config;
 use lmd_ghost::{LmdGhost, ThreadSafeReducedTree};
 use operation_pool::OperationPool;
 use parking_lot::RwLock;
-use slog::{info, Logger};
+use slog::{error, info, Logger};
 use slot_clock::{SlotClock, TestingSlotClock};
 use std::marker::PhantomData;
 use std::sync::Arc;
@@ -88,6 +89,8 @@ pub struct BeaconChainBuilder<T: BeaconChainTypes> {
     eth1_chain: Option<Eth1Chain<T::Eth1Chain, T::EthSpec>>,
     event_handler: Option<T::EventHandler>,
     slot_clock: Option<T::SlotClock>,
+    persisted_beacon_chain: Option<PersistedBeaconChain<T>>,
+    head_tracker: Option<HeadTracker>,
     spec: ChainSpec,
     log: Option<Logger>,
 }
@@ -128,6 +131,8 @@ where
             eth1_chain: None,
             event_handler: None,
             slot_clock: None,
+            persisted_beacon_chain: None,
+            head_tracker: None,
             spec: TEthSpec::default_spec(),
             log: None,
         }
@@ -214,6 +219,10 @@ where
         self.finalized_checkpoint = Some(p.canonical_head);
         self.genesis_block_root = Some(p.genesis_block_root);
 
+        self.head_tracker = HeadTracker::from_ssz_container(&p.ssz_head_tracker)
+            .map_err(|e| error!(log, "Failed to decode head tracker for database: {:?}", e))
+            .ok();
+
         Ok(self)
     }
 
@@ -379,6 +388,9 @@ where
             event_handler: self
                 .event_handler
                 .ok_or_else(|| "Cannot build without an event handler".to_string())?,
+            head_tracker: self
+                .head_tracker
+                .ok_or_else(|| "Cannot build without a head tracker".to_string())?,
             log: log.clone(),
         };
 

beacon_node/beacon_chain/src/head_tracker.rs

@@ -0,0 +1,187 @@
+use parking_lot::RwLock;
+use ssz_derive::{Decode, Encode};
+use std::collections::HashMap;
+use std::iter::FromIterator;
+use types::{BeaconBlock, EthSpec, Hash256, Slot};
+
+#[derive(Debug, PartialEq)]
+pub enum Error {
+    MismatchingLengths { roots_len: usize, slots_len: usize },
+}
+
+#[derive(Encode, Decode)]
+pub struct SszHeadTracker {
+    roots: Vec<Hash256>,
+    slots: Vec<Slot>,
+}
+
+#[derive(Default, Debug)]
+pub struct HeadTracker(RwLock<HashMap<Hash256, Slot>>);
+
+impl HeadTracker {
+    pub fn register_block<E: EthSpec>(&self, block_root: Hash256, block: &BeaconBlock<E>) {
+        let mut map = self.0.write();
+
+        map.remove(&block.parent_root);
+        map.insert(block_root, block.slot);
+    }
+
+    pub fn heads(&self) -> Vec<(Hash256, Slot)> {
+        self.0
+            .read()
+            .iter()
+            .map(|(root, slot)| (*root, *slot))
+            .collect()
+    }
+
+    pub fn to_ssz_container(&self) -> SszHeadTracker {
+        let (roots, slots) = self
+            .0
+            .read()
+            .iter()
+            .map(|(hash, slot)| (*hash, *slot))
+            .unzip();
+
+        SszHeadTracker { roots, slots }
+    }
+
+    pub fn from_ssz_container(ssz_container: &SszHeadTracker) -> Result<Self, Error> {
+        let roots_len = ssz_container.roots.len();
+        let slots_len = ssz_container.slots.len();
+
+        if roots_len != slots_len {
+            return Err(Error::MismatchingLengths {
+                roots_len,
+                slots_len,
+            });
+        } else {
+            let map = HashMap::from_iter(
+                ssz_container
+                    .roots
+                    .iter()
+                    .zip(ssz_container.slots.iter())
+                    .map(|(root, slot)| (*root, *slot)),
+            );
+
+            Ok(Self(RwLock::new(map)))
+        }
+    }
+}
+
+impl PartialEq<HeadTracker> for HeadTracker {
+    fn eq(&self, other: &HeadTracker) -> bool {
+        *self.0.read() == *other.0.read()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use ssz::{Decode, Encode};
+    use types::MainnetEthSpec;
+
+    type E = MainnetEthSpec;
+
+    #[test]
+    fn block_add() {
+        let spec = &E::default_spec();
+
+        let head_tracker = HeadTracker::default();
+
+        for i in 0..16 {
+            let mut block = BeaconBlock::empty(spec);
+            let block_root = Hash256::from_low_u64_be(i);
+
+            block.slot = Slot::new(i);
+            block.parent_root = if i == 0 {
+                Hash256::random()
+            } else {
+                Hash256::from_low_u64_be(i - 1)
+            };
+
+            head_tracker.register_block::<E>(block_root, &block);
+        }
+
+        assert_eq!(
+            head_tracker.heads(),
+            vec![(Hash256::from_low_u64_be(15), Slot::new(15))],
+            "should only have one head"
+        );
+
+        let mut block = BeaconBlock::empty(spec);
+        let block_root = Hash256::from_low_u64_be(42);
+        block.slot = Slot::new(15);
+        block.parent_root = Hash256::from_low_u64_be(14);
+        head_tracker.register_block::<E>(block_root, &block);
+
+        let heads = head_tracker.heads();
+
+        assert_eq!(heads.len(), 2, "should only have two heads");
+        assert!(
+            heads
+                .iter()
+                .any(|(root, slot)| *root == Hash256::from_low_u64_be(15) && *slot == Slot::new(15)),
+            "should contain first head"
+        );
+        assert!(
+            heads
+                .iter()
+                .any(|(root, slot)| *root == Hash256::from_low_u64_be(42) && *slot == Slot::new(15)),
+            "should contain second head"
+        );
+    }
+
+    #[test]
+    fn empty_round_trip() {
+        let non_empty = HeadTracker::default();
+        for i in 0..16 {
+            non_empty.0.write().insert(Hash256::random(), Slot::new(i));
+        }
+        let bytes = non_empty.to_ssz_container().as_ssz_bytes();
+
+        assert_eq!(
+            HeadTracker::from_ssz_container(
+                &SszHeadTracker::from_ssz_bytes(&bytes).expect("should decode")
+            ),
+            Ok(non_empty),
+            "non_empty should pass round trip"
+        );
+    }
+
+    #[test]
+    fn non_empty_round_trip() {
+        let non_empty = HeadTracker::default();
+        for i in 0..16 {
+            non_empty.0.write().insert(Hash256::random(), Slot::new(i));
+        }
+        let bytes = non_empty.to_ssz_container().as_ssz_bytes();
+
+        assert_eq!(
+            HeadTracker::from_ssz_container(
+                &SszHeadTracker::from_ssz_bytes(&bytes).expect("should decode")
+            ),
+            Ok(non_empty),
+            "non_empty should pass round trip"
+        );
+    }
+
+    #[test]
+    fn bad_length() {
+        let container = SszHeadTracker {
+            roots: vec![Hash256::random()],
+            slots: vec![],
+        };
+        let bytes = container.as_ssz_bytes();
+
+        assert_eq!(
+            HeadTracker::from_ssz_container(
+                &SszHeadTracker::from_ssz_bytes(&bytes).expect("should decode")
+            ),
+            Err(Error::MismatchingLengths {
+                roots_len: 1,
+                slots_len: 0
+            }),
+            "should fail decoding with bad lengths"
+        );
+    }
+}

beacon_node/beacon_chain/src/lib.rs

@@ -9,6 +9,7 @@ mod errors;
 pub mod eth1_chain;
 pub mod events;
 mod fork_choice;
+mod head_tracker;
 mod metrics;
 mod persisted_beacon_chain;
 pub mod test_utils;

beacon_node/beacon_chain/src/persisted_beacon_chain.rs

@@ -1,3 +1,4 @@
+use crate::head_tracker::SszHeadTracker;
 use crate::{BeaconChainTypes, CheckPoint};
 use operation_pool::PersistedOperationPool;
 use ssz::{Decode, Encode};
@@ -13,6 +14,7 @@ pub struct PersistedBeaconChain<T: BeaconChainTypes> {
     pub canonical_head: CheckPoint<T::EthSpec>,
     pub op_pool: PersistedOperationPool<T::EthSpec>,
     pub genesis_block_root: Hash256,
+    pub ssz_head_tracker: SszHeadTracker,
 }
 
 impl<T: BeaconChainTypes> SimpleStoreItem for PersistedBeaconChain<T> {