diff --git a/crates/net/p2p/src/req_resp/handlers.rs b/crates/net/p2p/src/req_resp/handlers.rs index 2a88a9bb..d2c63543 100644 --- a/crates/net/p2p/src/req_resp/handlers.rs +++ b/crates/net/p2p/src/req_resp/handlers.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use ethlambda_storage::Store; use libp2p::{PeerId, request_response}; @@ -269,29 +269,10 @@ fn canonical_blocks_by_range(store: &Store, start_slot: u64, count: u64) -> Vec< return Vec::new(); }; - let mut roots_by_slot = HashMap::new(); - let mut current_root = store.head(); - - while !current_root.is_zero() { - let Some(header) = store.get_block_header(¤t_root) else { - break; - }; - - if header.slot < start_slot { - break; - } - - if header.slot <= end_slot { - roots_by_slot.insert(header.slot, current_root); - } - - current_root = header.parent_root; - } - (start_slot..=end_slot) .filter_map(|slot| { - let root = roots_by_slot.get(&slot)?; - store.get_signed_block(root) + let root = store.get_block_root_by_slot(slot)?; + store.get_signed_block(&root) }) .collect() } diff --git a/crates/storage/src/api/tables.rs b/crates/storage/src/api/tables.rs index dcda1cbf..300b5293 100644 --- a/crates/storage/src/api/tables.rs +++ b/crates/storage/src/api/tables.rs @@ -10,6 +10,8 @@ pub enum Table { /// Stored separately from blocks because the genesis block has no signatures. /// All other blocks must have an entry in this table. BlockSignatures, + /// Canonical block index: slot -> block root + BlockRoots, /// State storage: H256 -> State /// /// Holds full-state snapshots only: the bootstrap anchor plus one anchor per @@ -32,10 +34,11 @@ pub enum Table { } /// All table variants. -pub const ALL_TABLES: [Table; 7] = [ +pub const ALL_TABLES: [Table; 8] = [ Table::BlockHeaders, Table::BlockBodies, Table::BlockSignatures, + Table::BlockRoots, Table::States, Table::StateDiffs, Table::Metadata, @@ -49,6 +52,7 @@ impl Table { Table::BlockHeaders => "block_headers", Table::BlockBodies => "block_bodies", Table::BlockSignatures => "block_signatures", + Table::BlockRoots => "block_roots", Table::States => "states", Table::StateDiffs => "state_diffs", Table::Metadata => "metadata", diff --git a/crates/storage/src/store.rs b/crates/storage/src/store.rs index c916bfca..524b9417 100644 --- a/crates/storage/src/store.rs +++ b/crates/storage/src/store.rs @@ -337,6 +337,10 @@ struct GossipDataEntry { /// Gossip signatures snapshot: (hashed_attestation_data, Vec<(validator_id, signature)>). pub type GossipSignatureSnapshot = Vec<(HashedAttestationData, Vec<(u64, ValidatorSignature)>)>; +type StorageKey = Vec; +type StorageEntry = (StorageKey, Vec); +type BlockRootIndexChanges = (Vec, Vec); + /// Bounded buffer for gossip signatures with FIFO eviction. /// /// Groups signatures by attestation data (via data_root). Each distinct @@ -518,12 +522,17 @@ fn decode_slot_root_key(bytes: &[u8]) -> (u64, H256) { (slot, root) } +fn encode_block_root_key(slot: u64) -> Vec { + slot.to_be_bytes().to_vec() +} + /// Fork choice store backed by a pluggable storage backend. /// /// The Store maintains all state required for fork choice and block processing: /// /// - **Metadata**: time, config, head, safe_target, justified/finalized checkpoints /// - **Blocks**: headers and bodies stored separately for efficient header-only queries +/// - **BlockRoots**: canonical block roots indexed by slot /// - **States**: beacon states indexed by block root /// - **Attestations**: latest known and pending ("new") attestations per validator /// - **Signatures**: gossip signatures and aggregated proofs for signature verification @@ -613,8 +622,7 @@ impl Store { ); return None; } - info!("Loaded store from persisted DB state"); - Some(Self { + let store = Self { backend, new_payloads: Arc::new(Mutex::new(PayloadBuffer::new(NEW_PAYLOAD_CAP))), known_payloads: Arc::new(Mutex::new(PayloadBuffer::new(AGGREGATED_PAYLOAD_CAP))), @@ -622,7 +630,12 @@ impl Store { GOSSIP_SIGNATURE_CAP, ))), state_cache: new_state_cache(), - }) + }; + if store.get_block_root_by_slot(store.head_slot()) != Some(store.head()) { + store.rebuild_block_root_index(); + } + info!("Loaded store from persisted DB state"); + Some(store) } /// Internal helper to initialize the store with anchor data. @@ -684,6 +697,16 @@ impl Store { .put_batch(Table::BlockHeaders, header_entries) .expect("put block header"); + batch + .put_batch( + Table::BlockRoots, + vec![( + encode_block_root_key(anchor_state.latest_block_header.slot), + anchor_block_root.to_ssz(), + )], + ) + .expect("put block root index"); + // Block body (if provided) if let Some(body) = anchor_body { let body_entries = vec![(anchor_block_root.to_ssz(), body.to_ssz())]; @@ -811,6 +834,9 @@ impl Store { pub fn update_checkpoints(&mut self, checkpoints: ForkCheckpoints) -> Result<(), Error> { // Read old finalized slot before updating metadata let old_finalized_slot = self.latest_finalized().slot; + let old_head = self.head(); + let (block_root_deletes, block_root_entries) = + self.block_root_index_changes(old_head, checkpoints.head); let mut entries = vec![(KEY_HEAD.to_vec(), checkpoints.head.to_ssz())]; @@ -824,6 +850,12 @@ impl Store { let mut batch = self.backend.begin_write().expect("write batch"); batch.put_batch(Table::Metadata, entries).expect("put"); + batch + .delete_batch(Table::BlockRoots, block_root_deletes) + .expect("delete old canonical block roots"); + batch + .put_batch(Table::BlockRoots, block_root_entries) + .expect("put canonical block roots"); batch.commit().expect("commit"); // Lightweight pruning that should happen immediately on finalization advance: @@ -873,6 +905,98 @@ impl Store { // ============ Blocks ============ + fn block_root_index_changes( + &self, + mut old_root: H256, + mut new_root: H256, + ) -> BlockRootIndexChanges { + let mut deletes = Vec::new(); + let mut entries = Vec::new(); + + while old_root != new_root { + if old_root.is_zero() { + let header = self + .get_block_header(&new_root) + .expect("new canonical block header exists"); + entries.push((encode_block_root_key(header.slot), new_root.to_ssz())); + new_root = header.parent_root; + continue; + } + if new_root.is_zero() { + let header = self + .get_block_header(&old_root) + .expect("old canonical block header exists"); + deletes.push(encode_block_root_key(header.slot)); + old_root = header.parent_root; + continue; + } + + let old_header = self + .get_block_header(&old_root) + .expect("old canonical block header exists"); + let new_header = self + .get_block_header(&new_root) + .expect("new canonical block header exists"); + + match old_header.slot.cmp(&new_header.slot) { + std::cmp::Ordering::Greater => { + deletes.push(encode_block_root_key(old_header.slot)); + old_root = old_header.parent_root; + } + std::cmp::Ordering::Less => { + entries.push((encode_block_root_key(new_header.slot), new_root.to_ssz())); + new_root = new_header.parent_root; + } + std::cmp::Ordering::Equal => { + deletes.push(encode_block_root_key(old_header.slot)); + entries.push((encode_block_root_key(new_header.slot), new_root.to_ssz())); + old_root = old_header.parent_root; + new_root = new_header.parent_root; + } + } + } + + (deletes, entries) + } + + fn rebuild_block_root_index(&self) { + let view = self.backend.begin_read().expect("read view"); + let old_keys = view + .prefix_iterator(Table::BlockRoots, &[]) + .expect("iterator") + .filter_map(Result::ok) + .map(|(key, _)| key.to_vec()) + .collect(); + drop(view); + + let mut entries = Vec::new(); + let mut root = self.head(); + while !root.is_zero() { + let Some(header) = self.get_block_header(&root) else { + break; + }; + entries.push((encode_block_root_key(header.slot), root.to_ssz())); + root = header.parent_root; + } + + let mut batch = self.backend.begin_write().expect("write batch"); + batch + .delete_batch(Table::BlockRoots, old_keys) + .expect("clear block root index"); + batch + .put_batch(Table::BlockRoots, entries) + .expect("rebuild block root index"); + batch.commit().expect("commit"); + } + + /// Return the canonical block root at `slot`. + pub fn get_block_root_by_slot(&self, slot: u64) -> Option { + let view = self.backend.begin_read().expect("read view"); + view.get(Table::BlockRoots, &encode_block_root_key(slot)) + .expect("get block root") + .map(|bytes| H256::from_ssz_bytes(&bytes).expect("valid block root")) + } + /// Get block data for fork choice: root -> (slot, parent_root). /// /// Iterates only the LiveChain table, avoiding Block deserialization. @@ -1576,6 +1700,12 @@ mod tests { vec![(encode_slot_root_key(slot, &root), vec![0u8; 4])], ) .expect("put sigs"); + batch + .put_batch( + Table::BlockRoots, + vec![(encode_block_root_key(slot), root.to_ssz())], + ) + .expect("put block root"); batch.commit().expect("commit"); } @@ -1618,6 +1748,19 @@ mod tests { H256::from(bytes) } + fn signed_block(slot: u64, parent_root: H256) -> SignedBlock { + SignedBlock { + message: Block { + slot, + proposer_index: 0, + parent_root, + state_root: H256::ZERO, + body: BlockBody::default(), + }, + proof: MultiMessageAggregate::default(), + } + } + impl Store { /// Create a Store with an in-memory backend for tests. fn test_store() -> Self { @@ -1651,7 +1794,88 @@ mod tests { // ============ Block Signature Pruning Tests ============ #[test] - fn prune_signatures_keeps_recent_window_when_finality_healthy() { + fn block_root_index_tracks_canonical_chain_across_reorgs() { + let backend = Arc::new(InMemoryBackend::new()); + let mut store = Store::from_anchor_state(backend, State::from_genesis(0, vec![])); + let anchor_root = store.head(); + + let block_1 = signed_block(1, anchor_root); + let root_1 = block_1.message.hash_tree_root(); + store + .insert_signed_block(root_1, block_1) + .expect("insert block 1"); + + let block_3 = signed_block(3, root_1); + let root_3 = block_3.message.hash_tree_root(); + store + .insert_signed_block(root_3, block_3) + .expect("insert block 3"); + store + .update_checkpoints(ForkCheckpoints::head_only(root_3)) + .expect("update head to block 3"); + + assert_eq!(store.get_block_root_by_slot(0), Some(anchor_root)); + assert_eq!(store.get_block_root_by_slot(1), Some(root_1)); + assert_eq!(store.get_block_root_by_slot(2), None); + assert_eq!(store.get_block_root_by_slot(3), Some(root_3)); + + let side_block_2 = signed_block(2, anchor_root); + let side_root_2 = side_block_2.message.hash_tree_root(); + store + .insert_signed_block(side_root_2, side_block_2) + .expect("insert side block 2"); + + let side_block_4 = signed_block(4, side_root_2); + let side_root_4 = side_block_4.message.hash_tree_root(); + store + .insert_signed_block(side_root_4, side_block_4) + .expect("insert side block 4"); + store + .update_checkpoints(ForkCheckpoints::head_only(side_root_4)) + .expect("update head to side block 4"); + + assert_eq!(store.get_block_root_by_slot(0), Some(anchor_root)); + assert_eq!(store.get_block_root_by_slot(1), None); + assert_eq!(store.get_block_root_by_slot(2), Some(side_root_2)); + assert_eq!(store.get_block_root_by_slot(3), None); + assert_eq!(store.get_block_root_by_slot(4), Some(side_root_4)); + } + + #[test] + fn from_db_state_rebuilds_block_root_index() { + let backend = Arc::new(InMemoryBackend::new()); + let mut store = + Store::from_anchor_state(backend.clone(), State::from_genesis(12345, vec![])); + + let block = signed_block(1, store.head()); + let block_root = block.message.hash_tree_root(); + store + .insert_signed_block(block_root, block) + .expect("insert block"); + store + .update_checkpoints(ForkCheckpoints::head_only(block_root)) + .expect("update head"); + + let view = backend.begin_read().expect("read view"); + let keys = view + .prefix_iterator(Table::BlockRoots, &[]) + .expect("iterator") + .filter_map(Result::ok) + .map(|(key, _)| key.to_vec()) + .collect(); + drop(view); + let mut batch = backend.begin_write().expect("write batch"); + batch + .delete_batch(Table::BlockRoots, keys) + .expect("clear block roots"); + batch.commit().expect("commit"); + + let restored = Store::from_db_state(backend, 12345).expect("restore store"); + assert_eq!(restored.get_block_root_by_slot(1), Some(block_root)); + } + + #[test] + fn prune_old_blocks_within_retention() { let backend = Arc::new(InMemoryBackend::new()); let mut store = Store::test_store_with_backend(backend.clone()); @@ -1670,6 +1894,9 @@ mod tests { // cutoff = 10: slots 0..9 pruned, slots 10..12 kept (within the window). assert_eq!(pruned, 10); + assert_eq!(count_entries(backend.as_ref(), Table::BlockSignatures), 3); + + // Oldest signatures are gone, but headers, bodies, and roots stay queryable. for i in 0..10u64 { assert!(!has_signature(backend.as_ref(), i, &root(i))); } @@ -1680,6 +1907,7 @@ mod tests { // Headers and bodies are always retained for the whole history. assert_eq!(count_entries(backend.as_ref(), Table::BlockHeaders), 13); assert_eq!(count_entries(backend.as_ref(), Table::BlockBodies), 13); + assert_eq!(count_entries(backend.as_ref(), Table::BlockRoots), 13); } #[test]