Arc-ify BeaconSnapshot

michaelsproul · Jun 28, 2022 · ccdc74d · ccdc74d
1 parent 67ca5d7
commit ccdc74d
Show file tree

Hide file tree

Showing 21 changed files with 493 additions and 232 deletions.
diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -1015,7 +1015,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     /// Returns the beacon block at the head of the canonical chain.
     ///
     /// See `Self::head` for more information.
-    pub fn head_beacon_block(&self) -> Result<SignedBeaconBlock<T::EthSpec>, Error> {
+    pub fn head_beacon_block(&self) -> Result<Arc<SignedBeaconBlock<T::EthSpec>>, Error> {
         self.with_head(|s| Ok(s.beacon_block.clone()))
     }
 
@@ -2964,7 +2964,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 snapshot_cache.insert(
                     BeaconSnapshot {
                         beacon_state: state,
-                        beacon_block: signed_block.as_ref().clone(),
+                        beacon_block: signed_block.clone(),
                         beacon_block_root: block_root,
                     },
                     None,
@@ -4446,10 +4446,13 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     ) -> Result<Vec<BeaconSnapshot<T::EthSpec, BlindedPayload<T::EthSpec>>>, Error> {
         let mut dump = vec![];
 
-        let mut last_slot = BeaconSnapshot {
-            beacon_block: self.head()?.beacon_block.into(),
-            beacon_block_root: self.head()?.beacon_block_root,
-            beacon_state: self.head()?.beacon_state,
+        let mut last_slot = {
+            let head = self.canonical_head.read();
+            BeaconSnapshot {
+                beacon_block: Arc::new(head.head_snapshot.beacon_block.clone_as_blinded()),
+                beacon_block_root: head.head_snapshot.beacon_block_root,
+                beacon_state: head.head_snapshot.beacon_state.clone(),
+            }
         };
 
         dump.push(last_slot.clone());
@@ -4476,7 +4479,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 })?;
 
             let slot = BeaconSnapshot {
-                beacon_block,
+                beacon_block: Arc::new(beacon_block),
                 beacon_block_root,
                 beacon_state,
             };

diff --git a/beacon_node/beacon_chain/src/beacon_snapshot.rs b/beacon_node/beacon_chain/src/beacon_snapshot.rs
@@ -1,4 +1,5 @@
 use serde_derive::Serialize;
+use std::sync::Arc;
 use types::{
     beacon_state::CloneConfig, BeaconState, EthSpec, ExecPayload, FullPayload, Hash256,
     SignedBeaconBlock,
@@ -8,15 +9,15 @@ use types::{
 /// head, justified head and finalized head.
 #[derive(Clone, Serialize, PartialEq, Debug)]
 pub struct BeaconSnapshot<E: EthSpec, Payload: ExecPayload<E> = FullPayload<E>> {
-    pub beacon_block: SignedBeaconBlock<E, Payload>,
+    pub beacon_block: Arc<SignedBeaconBlock<E, Payload>>,
     pub beacon_block_root: Hash256,
     pub beacon_state: BeaconState<E>,
 }
 
 impl<E: EthSpec, Payload: ExecPayload<E>> BeaconSnapshot<E, Payload> {
     /// Create a new checkpoint.
     pub fn new(
-        beacon_block: SignedBeaconBlock<E, Payload>,
+        beacon_block: Arc<SignedBeaconBlock<E, Payload>>,
         beacon_block_root: Hash256,
         beacon_state: BeaconState<E>,
     ) -> Self {
@@ -39,7 +40,7 @@ impl<E: EthSpec, Payload: ExecPayload<E>> BeaconSnapshot<E, Payload> {
     /// Update all fields of the checkpoint.
     pub fn update(
         &mut self,
-        beacon_block: SignedBeaconBlock<E, Payload>,
+        beacon_block: Arc<SignedBeaconBlock<E, Payload>>,
         beacon_block_root: Hash256,
         beacon_state: BeaconState<E>,
     ) {

diff --git a/beacon_node/beacon_chain/src/builder.rs b/beacon_node/beacon_chain/src/builder.rs
@@ -338,7 +338,7 @@ where
         Ok((
             BeaconSnapshot {
                 beacon_block_root,
-                beacon_block,
+                beacon_block: Arc::new(beacon_block),
                 beacon_state,
             },
             self,
@@ -459,7 +459,7 @@ where
 
         let snapshot = BeaconSnapshot {
             beacon_block_root: weak_subj_block_root,
-            beacon_block: weak_subj_block,
+            beacon_block: Arc::new(weak_subj_block),
             beacon_state: weak_subj_state,
         };
 
@@ -652,7 +652,7 @@ where
 
         let mut head_snapshot = BeaconSnapshot {
             beacon_block_root: head_block_root,
-            beacon_block: head_block,
+            beacon_block: Arc::new(head_block),
             beacon_state: head_state,
         };
 
@@ -728,12 +728,8 @@ where
         let genesis_validators_root = head_snapshot.beacon_state.genesis_validators_root();
         let genesis_time = head_snapshot.beacon_state.genesis_time();
         let head_for_snapshot_cache = head_snapshot.clone();
-        let fork_choice_view = fork_choice.cached_fork_choice_view();
-        let canonical_head = CanonicalHead {
-            fork_choice,
-            fork_choice_view,
-            head_snapshot,
-        };
+        let canonical_head = CanonicalHead::new(fork_choice, head_snapshot)
+            .map_err(|e| format!("Error creating canonical head: {:?}", e))?;
 
         let beacon_chain = BeaconChain {
             spec: self.spec,
@@ -1027,7 +1023,7 @@ mod test {
                 .get_blinded_block(&Hash256::zero())
                 .expect("should read db")
                 .expect("should find genesis block"),
-            block.clone().into(),
+            block.clone_as_blinded(),
             "should store genesis block under zero hash alias"
         );
         assert_eq!(