From 1a10ae12287d4c3ac6c3066bf864c888ba44f647 Mon Sep 17 00:00:00 2001 From: KaydenML Date: Sun, 23 Aug 2026 18:16:25 -0600 Subject: [PATCH] feat: keep slot_index canonical across side-branch imports and reorgs --- crates/common/fork_choice/lean/src/store.rs | 103 +++++++++++++++++++- crates/storage/src/db/lean.rs | 29 +++++- crates/storage/src/tables/lean/block.rs | 12 --- 3 files changed, 130 insertions(+), 14 deletions(-) diff --git a/crates/common/fork_choice/lean/src/store.rs b/crates/common/fork_choice/lean/src/store.rs index e76f2275f..ee98d0987 100644 --- a/crates/common/fork_choice/lean/src/store.rs +++ b/crates/common/fork_choice/lean/src/store.rs @@ -578,6 +578,7 @@ impl Store { }; let latest_finalized_checkpoint = latest_finalized_provider.get()?; + let old_head = head_provider.get()?; let finalized_slot = latest_finalized_checkpoint.slot; let attestations = { let mut relevant_keys = Vec::new(); @@ -657,12 +658,66 @@ impl Store { }; *self.network_state.finalized_checkpoint.write() = final_finalized_checkpoint; - head_provider.insert(new_head)?; + self.update_canonical_head(old_head, new_head).await?; latest_finalized_provider.insert(final_finalized_checkpoint)?; Ok(()) } + /// Align the slot index with the canonical chain after a fork-choice head change. + /// + /// Imported blocks are retained in the block and children tables regardless of + /// whether they are canonical. Consequently the slot index must be updated from + /// the old and new heads, rather than from block-import order. + async fn update_canonical_head(&self, old_head: B256, new_head: B256) -> anyhow::Result<()> { + if old_head == new_head { + return Ok(()); + } + + let block_provider = { + let db = self.store.lock().await; + db.block_provider() + }; + let mut new_entries = HashMap::::new(); + let mut stale_slots = HashSet::::new(); + let mut old_root = old_head; + let mut new_root = new_head; + + while old_root != new_root { + let new_block = block_provider.get(new_root)?; + let old_block = block_provider.get(old_root)?; + + match (new_block, old_block) { + (Some(block), old) + if old + .as_ref() + .is_none_or(|old| block.block.slot >= old.block.slot) => + { + new_entries.insert(block.block.slot, new_root); + new_root = block.block.parent_root; + } + (_, Some(block)) => { + stale_slots.insert(block.block.slot); + old_root = block.block.parent_root; + } + _ => break, + } + } + + let mut new_slots: Vec<_> = new_entries.into_iter().collect(); + new_slots.sort_unstable_by_key(|(slot, _)| *slot); + let new_slot_numbers: HashSet<_> = new_slots.iter().map(|(slot, _)| *slot).collect(); + let mut stale_slots: Vec<_> = stale_slots.difference(&new_slot_numbers).copied().collect(); + stale_slots.sort_unstable(); + + self.store + .lock() + .await + .update_canonical_head(new_head, &new_slots, &stale_slots)?; + + Ok(()) + } + pub async fn get_attestation_target(&self) -> anyhow::Result { let (head_provider, block_provider, safe_target_provider, state_provider) = { let db = self.store.lock().await; @@ -2786,6 +2841,52 @@ mod tests { ) } + #[tokio::test] + async fn side_branch_import_does_not_overwrite_slot_index() { + let store = sample_store_as_store(10).await; + let db = store.store.lock().await; + let genesis_root = db.head_provider().get().unwrap(); + let side_block = fake_signed_block(1, 1, genesis_root); + let side_root = side_block.block.tree_hash_root(); + + db.block_provider().insert(side_root, side_block).unwrap(); + + assert_eq!(db.slot_index_provider().get(1).unwrap(), None); + assert_eq!(db.head_provider().get().unwrap(), genesis_root); + } + + #[tokio::test] + async fn canonical_reindex_rewrites_and_deletes_reorged_slots() { + let store = sample_store_as_store(10).await; + let genesis_root = { store.store.lock().await.head_provider().get().unwrap() }; + let old_1 = fake_signed_block(1, 0, genesis_root); + let old_1_root = old_1.block.tree_hash_root(); + let old_2 = fake_signed_block(2, 0, old_1_root); + let old_2_root = old_2.block.tree_hash_root(); + let new_2 = fake_signed_block(2, 1, genesis_root); + let new_2_root = new_2.block.tree_hash_root(); + + { + let db = store.store.lock().await; + db.block_provider().insert(old_1_root, old_1).unwrap(); + db.block_provider().insert(old_2_root, old_2).unwrap(); + db.block_provider().insert(new_2_root, new_2).unwrap(); + db.slot_index_provider().insert(1, old_1_root).unwrap(); + db.slot_index_provider().insert(2, old_2_root).unwrap(); + db.head_provider().insert(old_2_root).unwrap(); + } + + store + .update_canonical_head(old_2_root, new_2_root) + .await + .unwrap(); + + let db = store.store.lock().await; + assert_eq!(db.slot_index_provider().get(1).unwrap(), None); + assert_eq!(db.slot_index_provider().get(2).unwrap(), Some(new_2_root)); + assert_eq!(db.head_provider().get().unwrap(), new_2_root); + } + #[tokio::test] async fn test_validate_attestation_rejects_head_on_finalized_orphaned_branch() { let (store, _, _, orphan_1_root, orphan_2_root) = diff --git a/crates/storage/src/db/lean.rs b/crates/storage/src/db/lean.rs index a7adcf68b..ded3bc2d8 100644 --- a/crates/storage/src/db/lean.rs +++ b/crates/storage/src/db/lean.rs @@ -1,6 +1,7 @@ use std::{fmt::Write, sync::Arc}; -use redb::{Database, ReadableDatabase, ReadableTableMetadata}; +use alloy_primitives::B256; +use redb::{Database, Durability, ReadableDatabase, ReadableTableMetadata}; use tracing::info; use crate::{ @@ -33,6 +34,32 @@ pub struct LeanDB { } impl LeanDB { + /// Atomically update the canonical slot index and persisted head. + pub fn update_canonical_head( + &self, + head: B256, + new_slots: &[(u64, B256)], + stale_slots: &[u64], + ) -> anyhow::Result<()> { + let mut write_txn = self.db.begin_write()?; + write_txn.set_durability(Durability::Immediate)?; + { + let mut slot_index = write_txn.open_table(LeanSlotIndexTable::TABLE_DEFINITION)?; + for (slot, root) in new_slots { + slot_index.insert(*slot, *root)?; + } + for slot in stale_slots { + slot_index.remove(*slot)?; + } + } + { + let mut head_field = write_txn.open_table(LeanHeadField::FIELD_DEFINITION)?; + head_field.insert(LeanHeadField::KEY, head)?; + } + write_txn.commit()?; + Ok(()) + } + /// Attach a cache to this LeanDB instance. /// This enables in-memory caching of blocks and states for improved performance. pub fn with_cache(mut self, cache: Arc) -> Self { diff --git a/crates/storage/src/tables/lean/block.rs b/crates/storage/src/tables/lean/block.rs index d9281fa4f..8128e60f3 100644 --- a/crates/storage/src/tables/lean/block.rs +++ b/crates/storage/src/tables/lean/block.rs @@ -7,7 +7,6 @@ use tree_hash::TreeHash; use super::{ children_index::{ChildIndexEntry, LeanChildrenIndexTable}, - slot_index::LeanSlotIndexTable, state_root_index::LeanStateRootIndexTable, }; use crate::{ @@ -82,10 +81,6 @@ impl REDBTable for LeanBlockTable { let mut table = write_txn.open_table(Self::TABLE_DEFINITION)?; let value = table.remove(key)?.map(|v| v.value()); if let Some(block) = &value { - let slot_index_table = LeanSlotIndexTable { - db: self.db.clone(), - }; - slot_index_table.remove(block.block.slot)?; let state_root_index_table = LeanStateRootIndexTable { db: self.db.clone(), }; @@ -109,14 +104,7 @@ impl LeanBlockTable { } pub fn insert_ref(&self, key: B256, value: &SignedBlock) -> Result<(), StoreError> { - // insert entry to slot_index table let block_root = value.block.tree_hash_root(); - let slot_index_table = LeanSlotIndexTable { - db: self.db.clone(), - }; - - slot_index_table.insert(value.block.slot, block_root)?; - // insert entry to state root index table let state_root_index_table = LeanStateRootIndexTable { db: self.db.clone(),