Query finalized root in storage too when performing peer status check

grandinetech · Jan 30, 2025 · dbae6e1 · dbae6e1
1 parent 846e07e
commit dbae6e1
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 16 deletions.
diff --git a/fork_choice_control/src/queries.rs b/fork_choice_control/src/queries.rs
@@ -158,12 +158,11 @@ where
         }
     }
 
-    // TODO(Grandine Team): This will incorrectly return `None` for archived slots.
-    #[must_use]
-    pub fn finalized_block_root_before_or_at(&self, slot: Slot) -> Option<H256> {
-        self.store_snapshot()
-            .finalized_before_or_at(slot)
-            .map(|chain_link| chain_link.block_root)
+    pub fn finalized_block_root_before_or_at(&self, slot: Slot) -> Result<Option<H256>> {
+        match self.store_snapshot().finalized_before_or_at(slot) {
+            Some(chain_link) => Ok(Some(chain_link.block_root)),
+            None => self.storage().block_root_before_or_at_slot(slot),
+        }
     }
 
     pub fn checkpoint_state(&self, checkpoint: Checkpoint) -> Result<Option<Arc<BeaconState<P>>>> {

diff --git a/fork_choice_control/src/storage.rs b/fork_choice_control/src/storage.rs
@@ -534,6 +534,21 @@ impl<P: Preset> Storage<P> {
         self.block_root_by_slot(slot)
     }
 
+    pub(crate) fn block_root_before_or_at_slot(&self, slot: Slot) -> Result<Option<H256>> {
+        let results = self
+            .database
+            .iterator_descending(..=BlockRootBySlot(slot).to_string())?;
+
+        itertools::process_results(results, |pairs| {
+            pairs
+                .take_while(|(key_bytes, _)| BlockRootBySlot::has_prefix(key_bytes))
+                .map(|(_, value_bytes)| H256::from_ssz_default(value_bytes))
+                .next()
+                .transpose()
+        })?
+        .map_err(Into::into)
+    }
+
     pub(crate) fn finalized_block_by_slot(
         &self,
         slot: Slot,
@@ -798,10 +813,8 @@ impl<P: Preset> Storage<P> {
 
         itertools::process_results(results, |pairs| {
             pairs
-                .take_while(|(key_bytes, _)| {
-                    FinalizedBlockByRoot::has_prefix(key_bytes)
-                        && !UnfinalizedBlockByRoot::has_prefix(key_bytes)
-                })
+                .take_while(|(key_bytes, _)| FinalizedBlockByRoot::has_prefix(key_bytes))
+                .filter(|(key_bytes, _)| !UnfinalizedBlockByRoot::has_prefix(key_bytes))
                 .count()
         })
     }
@@ -1213,4 +1226,41 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_block_root_before_or_at_slot() -> Result<()> {
+        let database = Database::in_memory();
+
+        database.put_batch(vec![
+            serialize(BlockRootBySlot(2), H256::repeat_byte(2))?,
+            serialize(BlockRootBySlot(6), H256::repeat_byte(6))?,
+        ])?;
+
+        let storage = Storage::<Mainnet>::new(
+            Arc::new(Config::mainnet()),
+            database,
+            nonzero!(64_u64),
+            StorageMode::Standard,
+        );
+
+        assert_eq!(storage.block_root_before_or_at_slot(1)?, None);
+        assert_eq!(
+            storage.block_root_before_or_at_slot(2)?,
+            Some(H256::repeat_byte(2)),
+        );
+        assert_eq!(
+            storage.block_root_before_or_at_slot(3)?,
+            Some(H256::repeat_byte(2)),
+        );
+        assert_eq!(
+            storage.block_root_before_or_at_slot(6)?,
+            Some(H256::repeat_byte(6)),
+        );
+        assert_eq!(
+            storage.block_root_before_or_at_slot(9)?,
+            Some(H256::repeat_byte(6)),
+        );
+
+        Ok(())
+    }
 }
diff --git a/p2p/src/network.rs b/p2p/src/network.rs
@@ -1581,11 +1581,18 @@ impl<P: Preset> Network<P> {
                 Ordering::Greater => {
                     let remote_finalized_slot = Self::start_of_epoch(remote.finalized_epoch);
 
-                    (
-                        self.controller
-                            .finalized_block_root_before_or_at(remote_finalized_slot),
-                        SyncStatus::Behind { info },
-                    )
+                    let finalized_root_at_slot = match self
+                        .controller
+                        .finalized_block_root_before_or_at(remote_finalized_slot)
+                    {
+                        Ok(root) => root,
+                        Err(error) => {
+                            warn!("failed to query for finalized block root: {error:?}");
+                            None
+                        }
+                    };
+
+                    (finalized_root_at_slot, SyncStatus::Behind { info })
                 }
             };
 
@@ -1606,7 +1613,10 @@ impl<P: Preset> Network<P> {
 
                 return;
             }
-        } else if matches!(sync_status, SyncStatus::Behind { .. }) {
+        } else if matches!(sync_status, SyncStatus::Behind { .. })
+            && local.finalized_root != H256::zero()
+            && remote.finalized_root != H256::zero()
+        {
             debug!(
                 "disconnecting peer {peer_id} due to missing historical data \
                  required to validate finalized root {:?} at epoch {}",