From d87e9ac270f2dd4859064bb5e5ef88a06186b0a5 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Tue, 7 Oct 2025 16:54:53 -0400 Subject: [PATCH 01/14] wip --- crates/node/src/full/mod.rs | 2 ++ crates/sync/pipeline/src/lib.rs | 1 + crates/sync/stage/src/blocks/downloader.rs | 11 ++++++++--- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/crates/node/src/full/mod.rs b/crates/node/src/full/mod.rs index 599302d22..43afc339f 100644 --- a/crates/node/src/full/mod.rs +++ b/crates/node/src/full/mod.rs @@ -202,6 +202,8 @@ impl Node { } pub async fn launch(self) -> Result { + println!("Launching node"); + if let Some(ref cfg) = self.config.metrics { let reports: Vec> = vec![Box::new(self.db.clone()) as Box]; let exporter = PrometheusRecorder::current().expect("qed; should exist at this point"); diff --git a/crates/sync/pipeline/src/lib.rs b/crates/sync/pipeline/src/lib.rs index aeeb562e3..66a8803ef 100644 --- a/crates/sync/pipeline/src/lib.rs +++ b/crates/sync/pipeline/src/lib.rs @@ -65,6 +65,7 @@ //! [Erigon]: https://github.com/erigontech/erigon use core::future::IntoFuture; +use std::sync::atomic::{AtomicU64, Ordering}; use futures::future::BoxFuture; use katana_primitives::block::BlockNumber; diff --git a/crates/sync/stage/src/blocks/downloader.rs b/crates/sync/stage/src/blocks/downloader.rs index 5764e0079..11af2a174 100644 --- a/crates/sync/stage/src/blocks/downloader.rs +++ b/crates/sync/stage/src/blocks/downloader.rs @@ -16,6 +16,7 @@ use anyhow::Result; use katana_gateway::client::Client as GatewayClient; use katana_gateway::types::StateUpdateWithBlock; use katana_primitives::block::BlockNumber; +use tracing::{info, info_span, trace, Instrument}; use crate::downloader::{BatchDownloader, Downloader}; @@ -86,9 +87,12 @@ where to: BlockNumber, ) -> impl Future, katana_gateway::client::Error>> + Send { - // convert the range to a list of block keys - let block_keys = (from..=to).collect::>(); - self.inner.download(block_keys) + async move { + // convert the range to a list of block keys + let block_keys = (from..=to).collect::>(); + self.inner.download(block_keys).await + } + .instrument(info_span!("download_blocks", %from, %to)) } } @@ -125,6 +129,7 @@ mod impls { &self, key: &Self::Key, ) -> impl Future> { + info!(block = %key, "Downloading block."); async { match self.gateway.get_state_update_with_block((*key).into()).await.inspect_err( |error| error!(block = %*key, ?error, "Error downloading block from gateway."), From 602996b2341c3b124bc96a59dabf34e912cdb354 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Tue, 7 Oct 2025 23:44:38 -0400 Subject: [PATCH 02/14] wip --- crates/sync/pipeline/src/lib.rs | 5 +++++ crates/sync/stage/src/blocks/downloader.rs | 4 ++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/sync/pipeline/src/lib.rs b/crates/sync/pipeline/src/lib.rs index 66a8803ef..c49144a45 100644 --- a/crates/sync/pipeline/src/lib.rs +++ b/crates/sync/pipeline/src/lib.rs @@ -142,6 +142,11 @@ impl PipelineHandle { pub async fn stopped(&self) { self.tx.closed().await; } + + /// Wait until the [`Pipeline`] has stopped. + pub async fn stopped(&self) { + self.tx.closed().await; + } } /// Syncing pipeline. diff --git a/crates/sync/stage/src/blocks/downloader.rs b/crates/sync/stage/src/blocks/downloader.rs index 11af2a174..9bb036849 100644 --- a/crates/sync/stage/src/blocks/downloader.rs +++ b/crates/sync/stage/src/blocks/downloader.rs @@ -16,7 +16,7 @@ use anyhow::Result; use katana_gateway::client::Client as GatewayClient; use katana_gateway::types::StateUpdateWithBlock; use katana_primitives::block::BlockNumber; -use tracing::{info, info_span, trace, Instrument}; +use tracing::{info_span, Instrument}; use crate::downloader::{BatchDownloader, Downloader}; @@ -129,7 +129,7 @@ mod impls { &self, key: &Self::Key, ) -> impl Future> { - info!(block = %key, "Downloading block."); + trace!(block = %key, "Downloading block."); async { match self.gateway.get_state_update_with_block((*key).into()).await.inspect_err( |error| error!(block = %*key, ?error, "Error downloading block from gateway."), From de0994a0c4e0e2862126bc09b2b0370c5a20ec40 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Thu, 9 Oct 2025 16:35:08 -0400 Subject: [PATCH 03/14] wip --- crates/node/src/full/mod.rs | 2 -- crates/sync/stage/src/blocks/mod.rs | 3 ++- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/crates/node/src/full/mod.rs b/crates/node/src/full/mod.rs index 43afc339f..599302d22 100644 --- a/crates/node/src/full/mod.rs +++ b/crates/node/src/full/mod.rs @@ -202,8 +202,6 @@ impl Node { } pub async fn launch(self) -> Result { - println!("Launching node"); - if let Some(ref cfg) = self.config.metrics { let reports: Vec> = vec![Box::new(self.db.clone()) as Box]; let exporter = PrometheusRecorder::current().expect("qed; should exist at this point"); diff --git a/crates/sync/stage/src/blocks/mod.rs b/crates/sync/stage/src/blocks/mod.rs index 243388e6d..cb9afd9e4 100644 --- a/crates/sync/stage/src/blocks/mod.rs +++ b/crates/sync/stage/src/blocks/mod.rs @@ -106,7 +106,8 @@ where .download_blocks(input.from(), input.to()) .instrument(info_span!(target: "stage", "blocks.download", from = %input.from(), to = %input.to())) .await - .map_err(Error::Gateway)?; + .map_err(Error::Gateway) + .inspect_err(|e| error!(error = %e , "Error downloading blocks."))?; let span = info_span!(target: "stage", "blocks.insert", from = %input.from(), to = %input.to()); let _enter = span.enter(); From 84cf683b4a2f6617931300fd9d6ddca2f7dd8a37 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Mon, 13 Oct 2025 20:15:56 -0400 Subject: [PATCH 04/14] wip --- crates/cli/src/full.rs | 1 - crates/sync/pipeline/src/lib.rs | 3 ++- crates/sync/stage/src/blocks/downloader.rs | 11 +++-------- crates/sync/stage/src/blocks/mod.rs | 3 +-- crates/sync/stage/src/classes.rs | 1 + 5 files changed, 7 insertions(+), 12 deletions(-) diff --git a/crates/cli/src/full.rs b/crates/cli/src/full.rs index 1fb1ddcd3..623f6b848 100644 --- a/crates/cli/src/full.rs +++ b/crates/cli/src/full.rs @@ -2,7 +2,6 @@ use std::path::PathBuf; use anyhow::{Context, Result}; pub use clap::Parser; -use katana_node::config::db::DbConfig; use katana_node::config::metrics::MetricsConfig; use katana_node::config::rpc::RpcConfig; use katana_node::full; diff --git a/crates/sync/pipeline/src/lib.rs b/crates/sync/pipeline/src/lib.rs index c49144a45..fd1e71cb6 100644 --- a/crates/sync/pipeline/src/lib.rs +++ b/crates/sync/pipeline/src/lib.rs @@ -65,7 +65,6 @@ //! [Erigon]: https://github.com/erigontech/erigon use core::future::IntoFuture; -use std::sync::atomic::{AtomicU64, Ordering}; use futures::future::BoxFuture; use katana_primitives::block::BlockNumber; @@ -291,6 +290,8 @@ impl Pipeline

{ /// Returns an error if any stage execution fails or if the pipeline fails to read the /// checkpoint. pub async fn run_once(&mut self, to: BlockNumber) -> PipelineResult { + let tip = self.tip.expect("qed; should exist by now"); + if self.stages.is_empty() { return Ok(to); } diff --git a/crates/sync/stage/src/blocks/downloader.rs b/crates/sync/stage/src/blocks/downloader.rs index 9bb036849..5764e0079 100644 --- a/crates/sync/stage/src/blocks/downloader.rs +++ b/crates/sync/stage/src/blocks/downloader.rs @@ -16,7 +16,6 @@ use anyhow::Result; use katana_gateway::client::Client as GatewayClient; use katana_gateway::types::StateUpdateWithBlock; use katana_primitives::block::BlockNumber; -use tracing::{info_span, Instrument}; use crate::downloader::{BatchDownloader, Downloader}; @@ -87,12 +86,9 @@ where to: BlockNumber, ) -> impl Future, katana_gateway::client::Error>> + Send { - async move { - // convert the range to a list of block keys - let block_keys = (from..=to).collect::>(); - self.inner.download(block_keys).await - } - .instrument(info_span!("download_blocks", %from, %to)) + // convert the range to a list of block keys + let block_keys = (from..=to).collect::>(); + self.inner.download(block_keys) } } @@ -129,7 +125,6 @@ mod impls { &self, key: &Self::Key, ) -> impl Future> { - trace!(block = %key, "Downloading block."); async { match self.gateway.get_state_update_with_block((*key).into()).await.inspect_err( |error| error!(block = %*key, ?error, "Error downloading block from gateway."), diff --git a/crates/sync/stage/src/blocks/mod.rs b/crates/sync/stage/src/blocks/mod.rs index cb9afd9e4..243388e6d 100644 --- a/crates/sync/stage/src/blocks/mod.rs +++ b/crates/sync/stage/src/blocks/mod.rs @@ -106,8 +106,7 @@ where .download_blocks(input.from(), input.to()) .instrument(info_span!(target: "stage", "blocks.download", from = %input.from(), to = %input.to())) .await - .map_err(Error::Gateway) - .inspect_err(|e| error!(error = %e , "Error downloading blocks."))?; + .map_err(Error::Gateway)?; let span = info_span!(target: "stage", "blocks.insert", from = %input.from(), to = %input.to()); let _enter = span.enter(); diff --git a/crates/sync/stage/src/classes.rs b/crates/sync/stage/src/classes.rs index 955409102..f5b5dae30 100644 --- a/crates/sync/stage/src/classes.rs +++ b/crates/sync/stage/src/classes.rs @@ -153,6 +153,7 @@ where for (key, class) in declared_class_hashes.iter().zip(verified_classes.into_iter()) { self.provider.set_class(key.class_hash, class)?; } + } else { } Ok(StageExecutionOutput { last_block_processed: input.to() }) From e91eaa573d01884e99271e19cc437b347996e992 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Mon, 13 Oct 2025 21:52:59 -0400 Subject: [PATCH 05/14] feat(sync): add unwinding support for chain reorganizations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement functionality to handle chain reorganizations by unwinding blocks in the sync pipeline. When the pipeline receives a new tip that is smaller than its current tip, it can now revert the chain state back to a specific block number. The Stage trait has been extended with a new unwind method that accepts a target block number and removes all blocks after it, ensuring the database state reflects only the data up to the unwind point. Each stage implements unwinding differently based on its responsibilities. The Blocks stage removes block headers, hashes, transactions, and receipts from their respective tables. The Classes stage removes class declarations and artifacts for blocks after the unwind target. The StateTrie stage unwinding is left unimplemented as it requires more complex trie unwinding logic that maintains merkle tree invariants. All stages update their checkpoints after unwinding to maintain consistency. The implementation uses the existing database transaction abstraction to ensure atomicity of the unwind operations. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- Cargo.lock | 1 + crates/sync/stage/Cargo.toml | 1 + crates/sync/stage/src/blocks/mod.rs | 115 +++++++++++++++++++++++++++- crates/sync/stage/src/classes.rs | 62 ++++++++++++++- crates/sync/stage/src/lib.rs | 25 +++++- crates/sync/stage/src/trie.rs | 15 +++- 6 files changed, 214 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a1fcc96cc..045ae6df3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6597,6 +6597,7 @@ dependencies = [ "backon", "futures", "katana-core", + "katana-db", "katana-executor", "katana-gateway", "katana-messaging", diff --git a/crates/sync/stage/Cargo.toml b/crates/sync/stage/Cargo.toml index a04371839..f5caca0ad 100644 --- a/crates/sync/stage/Cargo.toml +++ b/crates/sync/stage/Cargo.toml @@ -7,6 +7,7 @@ version.workspace = true [dependencies] katana-core.workspace = true +katana-db.workspace = true katana-executor.workspace = true katana-gateway.workspace = true katana-messaging.workspace = true diff --git a/crates/sync/stage/src/blocks/mod.rs b/crates/sync/stage/src/blocks/mod.rs index 243388e6d..e2ef038df 100644 --- a/crates/sync/stage/src/blocks/mod.rs +++ b/crates/sync/stage/src/blocks/mod.rs @@ -1,8 +1,10 @@ use anyhow::Result; use futures::future::BoxFuture; +use katana_db::abstraction::{Database, DbCursor, DbTx, DbTxMut}; +use katana_db::tables; use katana_gateway::types::{BlockStatus, StateUpdate as GatewayStateUpdate, StateUpdateWithBlock}; use katana_primitives::block::{ - FinalityStatus, GasPrices, Header, SealedBlock, SealedBlockWithStatus, + BlockNumber, FinalityStatus, GasPrices, Header, SealedBlock, SealedBlockWithStatus, }; use katana_primitives::fee::{FeeInfo, PriceUnit}; use katana_primitives::receipt::{ @@ -13,6 +15,7 @@ use katana_primitives::transaction::{Tx, TxWithHash}; use katana_primitives::Felt; use katana_provider::api::block::{BlockHashProvider, BlockWriter}; use katana_provider::ProviderError; +use katana_provider::api::stage::StageCheckpointProvider; use num_traits::ToPrimitive; use starknet::core::types::ResourcePrice; use tracing::{error, info_span, Instrument}; @@ -88,11 +91,105 @@ impl Blocks { Ok(()) } + + /// Unwinds block data by removing all blocks after the specified block number. + /// + /// This removes entries from the following tables: + /// - Headers, BlockHashes, BlockNumbers, BlockBodyIndices, BlockStatusses + /// - TxNumbers, TxBlocks, TxHashes, TxTraces, Transactions, Receipts + fn unwind_blocks(db: &Db, unwind_to: BlockNumber) -> Result<(), crate::Error> { + db.update(|db_tx| -> Result<(), katana_provider::api::ProviderError> { + // Get the tx_offset for the unwind_to block to know where to start deleting txs + let mut last_tx_num = None; + if let Some(indices) = db_tx.get::(unwind_to)? { + last_tx_num = Some(indices.tx_offset + indices.tx_count); + } + + // Remove all blocks after unwind_to + let mut blocks_to_remove = Vec::new(); + let mut cursor = db_tx.cursor_mut::()?; + + // Find all blocks after unwind_to + if let Some((block_num, _)) = cursor.seek(unwind_to + 1)? { + blocks_to_remove.push(block_num); + while let Some((block_num, _)) = cursor.next()? { + blocks_to_remove.push(block_num); + } + } + drop(cursor); + + // Remove block data + for block_num in blocks_to_remove { + // Get block hash before deleting + let block_hash = db_tx.get::(block_num)?; + + db_tx.delete::(block_num, None)?; + db_tx.delete::(block_num, None)?; + db_tx.delete::(block_num, None)?; + db_tx.delete::(block_num, None)?; + + if let Some(hash) = block_hash { + db_tx.delete::(hash, None)?; + } + } + + // Remove transaction data if we have a last_tx_num + if let Some(start_tx_num) = last_tx_num { + let mut txs_to_remove = Vec::new(); + let mut cursor = db_tx.cursor_mut::()?; + + if let Some((tx_num, _)) = cursor.seek(start_tx_num)? { + txs_to_remove.push(tx_num); + while let Some((tx_num, _)) = cursor.next()? { + txs_to_remove.push(tx_num); + } + } + drop(cursor); + + for tx_num in txs_to_remove { + // Get tx hash before deleting + let tx_hash = db_tx.get::(tx_num)?; + + db_tx.delete::(tx_num, None)?; + db_tx.delete::(tx_num, None)?; + db_tx.delete::(tx_num, None)?; + db_tx.delete::(tx_num, None)?; + db_tx.delete::(tx_num, None)?; + + if let Some(hash) = tx_hash { + db_tx.delete::(hash, None)?; + } + } + } + + Ok(()) + }) + .map_err(katana_provider::api::ProviderError::from)??; + + Ok(()) + } +} + +/// Trait for accessing the database from a provider. +pub trait DatabaseProvider { + /// The database type. + type Db: Database; + + /// Returns a reference to the underlying database. + fn db(&self) -> &Self::Db; +} + +impl DatabaseProvider for katana_provider::providers::db::DbProvider { + type Db = Db; + + fn db(&self) -> &Self::Db { + katana_provider::providers::db::DbProvider::db(self) + } } impl Stage for Blocks where - P: BlockWriter + BlockHashProvider, + P: BlockWriter + BlockHashProvider + DatabaseProvider + StageCheckpointProvider, D: BlockDownloader, { fn id(&self) -> &'static str { @@ -133,6 +230,20 @@ where Ok(StageExecutionOutput { last_block_processed: input.to() }) }) } + + fn unwind<'a>(&'a mut self, unwind_to: BlockNumber) -> BoxFuture<'a, StageResult> { + Box::pin(async move { + debug!(target: "stage", id = %self.id(), unwind_to = %unwind_to, "Unwinding blocks."); + + // Unwind blocks + Self::unwind_blocks(self.provider.db(), unwind_to)?; + + // Update checkpoint + self.provider.set_checkpoint(self.id(), unwind_to)?; + + Ok(StageExecutionOutput { last_block_processed: unwind_to }) + }) + } } #[derive(Debug, thiserror::Error)] diff --git a/crates/sync/stage/src/classes.rs b/crates/sync/stage/src/classes.rs index f5b5dae30..fc1b4d30c 100644 --- a/crates/sync/stage/src/classes.rs +++ b/crates/sync/stage/src/classes.rs @@ -3,11 +3,14 @@ use std::future::Future; use anyhow::Result; use futures::channel::oneshot; use futures::future::BoxFuture; +use katana_db::abstraction::{Database, DbCursor, DbTxMut}; +use katana_db::tables; use katana_gateway::client::Client as SequencerGateway; use katana_gateway::types::ContractClass as GatewayContractClass; use katana_primitives::block::BlockNumber; use katana_primitives::class::{ClassHash, ContractClass}; use katana_provider::api::contract::ContractClassWriter; +use katana_provider::api::stage::StageCheckpointProvider; use katana_provider::api::state_update::StateUpdateProvider; use katana_provider::api::ProviderError; use katana_rpc_types::class::ConversionError; @@ -15,6 +18,7 @@ use rayon::prelude::*; use tracing::{debug, error, info, info_span, Instrument}; use super::{Stage, StageExecutionInput, StageExecutionOutput, StageResult}; +use crate::blocks::DatabaseProvider; use crate::downloader::{BatchDownloader, Downloader, DownloaderResult}; /// A stage for downloading and storing contract classes. @@ -49,6 +53,48 @@ impl

Classes

{ Self { provider, downloader, verification_pool } } + /// Unwinds class data by removing all classes declared after the specified block number. + /// + /// This removes entries from the following tables: + /// - CompiledClassHashes, Classes, ClassDeclarationBlock, ClassDeclarations + fn unwind_classes(db: &Db, unwind_to: BlockNumber) -> Result<(), crate::Error> { + db.update(|db_tx| -> Result<(), katana_provider::api::ProviderError> { + // Find all classes declared after unwind_to + let mut classes_to_remove = Vec::new(); + let mut cursor = db_tx.cursor_dup_mut::()?; + + // Find all blocks after unwind_to that have class declarations + if let Some((block_num, class_hash)) = cursor.seek(unwind_to + 1)? { + classes_to_remove.push((block_num, class_hash)); + + while let Some((block_num, class_hash)) = cursor.next()? { + classes_to_remove.push((block_num, class_hash)); + } + } + drop(cursor); + + // Remove class declarations for blocks after unwind_to + for (block_num, class_hash) in &classes_to_remove { + // Delete from ClassDeclarations (dupsort table) + db_tx.delete::(*block_num, Some(*class_hash))?; + + // Delete from ClassDeclarationBlock + db_tx.delete::(*class_hash, None)?; + + // Delete the class itself from Classes + db_tx.delete::(*class_hash, None)?; + + // Delete compiled class hash + db_tx.delete::(*class_hash, None)?; + } + + Ok(()) + }) + .map_err(katana_provider::api::ProviderError::from)??; + + Ok(()) + } + /// Returns the hashes of the classes declared in the given range of blocks. fn get_declared_classes( &self, @@ -120,7 +166,7 @@ impl

Classes

{ impl

Stage for Classes

where - P: StateUpdateProvider + ContractClassWriter, + P: StateUpdateProvider + ContractClassWriter + DatabaseProvider + StageCheckpointProvider, { fn id(&self) -> &'static str { "Classes" @@ -159,6 +205,20 @@ where Ok(StageExecutionOutput { last_block_processed: input.to() }) }) } + + fn unwind<'a>(&'a mut self, unwind_to: BlockNumber) -> BoxFuture<'a, StageResult> { + Box::pin(async move { + debug!(target: "stage", id = %self.id(), unwind_to = %unwind_to, "Unwinding classes."); + + // Unwind classes + Self::unwind_classes(self.provider.db(), unwind_to)?; + + // Update checkpoint + self.provider.set_checkpoint(self.id(), unwind_to)?; + + Ok(StageExecutionOutput { last_block_processed: unwind_to }) + }) + } } #[derive(Debug, thiserror::Error)] diff --git a/crates/sync/stage/src/lib.rs b/crates/sync/stage/src/lib.rs index 3fbb05c58..10e413e2f 100644 --- a/crates/sync/stage/src/lib.rs +++ b/crates/sync/stage/src/lib.rs @@ -10,7 +10,7 @@ pub mod downloader; mod sequencing; pub mod trie; -pub use blocks::Blocks; +pub use blocks::{Blocks, DatabaseProvider}; pub use classes::Classes; pub use sequencing::Sequencing; pub use trie::StateTrie; @@ -126,6 +126,29 @@ pub trait Stage: Send + Sync { /// Implementors are expected to perform any necessary processings on all blocks in the range /// `[input.from, input.to]`. fn execute<'a>(&'a mut self, input: &'a StageExecutionInput) -> BoxFuture<'a, StageResult>; + + /// Unwinds the stage to the specified block number. + /// + /// This method is called during chain reorganizations to revert the chain state back to a + /// specific block. All blocks after the `unwind_to` block should be removed, and the + /// resulting database state should be as if the stage had only synced up to `unwind_to`. + /// + /// # Arguments + /// + /// * `unwind_to` - The target block number to unwind to. All blocks after this will be removed. + /// + /// # Returns + /// + /// A future that resolves to a [`StageResult`] containing [`StageExecutionOutput`] + /// with the last block number after unwinding (should equal `unwind_to`). + /// + /// # Implementation Requirements + /// + /// Implementors must ensure that: + /// - All data for blocks > `unwind_to` is removed from relevant database tables + /// - The stage checkpoint is updated to reflect the unwound state + /// - Database invariants are maintained after the unwind operation + fn unwind<'a>(&'a mut self, unwind_to: BlockNumber) -> BoxFuture<'a, StageResult>; } #[cfg(test)] diff --git a/crates/sync/stage/src/trie.rs b/crates/sync/stage/src/trie.rs index ee1f90752..bf6faa953 100644 --- a/crates/sync/stage/src/trie.rs +++ b/crates/sync/stage/src/trie.rs @@ -8,7 +8,7 @@ use katana_rpc_types::class; use katana_trie::CommitId; use starknet::macros::short_string; use starknet_types_core::hash::{Poseidon, StarkHash}; -use tracing::{debug, debug_span, error}; +use tracing::{debug, debug_span, error, warn}; use crate::{Stage, StageExecutionInput, StageExecutionOutput, StageResult}; @@ -109,6 +109,19 @@ where Ok(StageExecutionOutput { last_block_processed: input.to() }) }) } + + fn unwind<'a>(&'a mut self, unwind_to: BlockNumber) -> BoxFuture<'a, StageResult> { + Box::pin(async move { + warn!( + target: "stage", + id = %self.id(), + unwind_to = %unwind_to, + "StateTrie unwinding not implemented - requires complex trie unwinding logic" + ); + + unimplemented!("StateTrie unwinding requires complex trie unwinding logic") + }) + } } #[derive(Debug, thiserror::Error)] From e1ff3b54e76c50cf68af40018492be7032a83eb3 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Thu, 16 Oct 2025 20:43:23 -0400 Subject: [PATCH 06/14] wip --- Cargo.lock | 2 +- crates/storage/db/src/trie/mod.rs | 66 ++++++- .../storage/provider/provider-api/src/trie.rs | 8 + .../provider/src/providers/db/trie.rs | 46 ++++- crates/sync/pipeline/src/lib.rs | 178 +++++++++++++++--- crates/sync/pipeline/tests/pipeline.rs | 32 ++-- crates/sync/stage/src/blocks/mod.rs | 2 +- crates/sync/stage/src/lib.rs | 7 +- crates/sync/stage/src/trie.rs | 11 +- crates/trie/Cargo.toml | 2 +- crates/trie/src/classes.rs | 4 + crates/trie/src/contracts.rs | 4 + crates/trie/src/lib.rs | 58 +++++- crates/trie/src/storages.rs | 4 + 14 files changed, 355 insertions(+), 69 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 045ae6df3..b438df3f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2018,7 +2018,7 @@ dependencies = [ [[package]] name = "bonsai-trie" version = "0.1.0" -source = "git+https://github.com/dojoengine/bonsai-trie/?rev=8b509fd#8b509fdf2c22fb7386acc6c00f179ac1cd099243" +source = "git+https://github.com/dojoengine/bonsai-trie/?rev=95de4c6#95de4c600f20cec326650364dea841d23cf93718" dependencies = [ "bitvec", "derive_more 0.99.20", diff --git a/crates/storage/db/src/trie/mod.rs b/crates/storage/db/src/trie/mod.rs index 2f8242e67..aba812b3a 100644 --- a/crates/storage/db/src/trie/mod.rs +++ b/crates/storage/db/src/trie/mod.rs @@ -173,9 +173,22 @@ where fn get_by_prefix( &self, - _: &DatabaseKey<'_>, + prefix: &DatabaseKey<'_>, ) -> Result, Self::DatabaseError> { - todo!() + let mut results = Vec::new(); + + let mut cursor = self.tx.cursor::()?; + let walker = cursor.walk(None)?; + + for entry in walker { + let (TrieDatabaseKey { key, .. }, value) = entry?; + + if key.starts_with(prefix.as_slice()) { + results.push((key.to_smallvec(), value)); + } + } + + Ok(results) } fn insert( @@ -300,8 +313,7 @@ where &self, prefix: &DatabaseKey<'_>, ) -> Result, Self::DatabaseError> { - let _ = prefix; - todo!() + TrieDb::::new(self.tx.clone()).get_by_prefix(prefix) } fn insert( @@ -509,4 +521,50 @@ mod tests { assert_eq!(vec![value0, value1], result); } } + + #[test] + fn revert_to() { + let db = test_utils::create_test_db(); + let db_tx = db.tx_mut().expect("failed to get tx"); + + let mut trie = ClassesTrie::new(TrieDbMut::::new(&db_tx)); + + // Insert values at block 0 + trie.insert(felt!("0x1"), felt!("0x100")); + trie.insert(felt!("0x2"), felt!("0x200")); + trie.commit(0); + let root_at_block_0 = trie.root(); + + // Insert more values at block 1 + trie.insert(felt!("0x3"), felt!("0x300")); + trie.insert(felt!("0x4"), felt!("0x400")); + trie.commit(1); + let root_at_block_1 = trie.root(); + + // Roots should be different + assert_ne!(root_at_block_0, root_at_block_1); + + // Insert even more values at block 2 + trie.insert(felt!("0x5"), felt!("0x500")); + trie.commit(2); + let root_at_block_2 = trie.root(); + + // Roots should be different + assert_ne!(root_at_block_1, root_at_block_2); + assert_ne!(root_at_block_0, root_at_block_2); + + // Revert to block 1 + trie.revert_to(1, 2); + let root_after_revert = trie.root(); + + // After revert, root should match block 1 + assert_eq!(root_after_revert, root_at_block_1); + + // Revert to block 0 + trie.revert_to(0, 1); + let root_after_second_revert = trie.root(); + + // After revert, root should match block 0 + assert_eq!(root_after_second_revert, root_at_block_0); + } } diff --git a/crates/storage/provider/provider-api/src/trie.rs b/crates/storage/provider/provider-api/src/trie.rs index 6a63ec3c4..d570189c6 100644 --- a/crates/storage/provider/provider-api/src/trie.rs +++ b/crates/storage/provider/provider-api/src/trie.rs @@ -20,4 +20,12 @@ pub trait TrieWriter: Send + Sync { block_number: BlockNumber, state_updates: &StateUpdates, ) -> ProviderResult; + + fn unwind_classes_trie(&self, unwind_to: BlockNumber) -> ProviderResult { + unimplemented!() + } + + fn unwind_contracts_trie(&self, unwind_to: BlockNumber) -> ProviderResult { + unimplemented!() + } } diff --git a/crates/storage/provider/provider/src/providers/db/trie.rs b/crates/storage/provider/provider/src/providers/db/trie.rs index 912d510b6..6795232fd 100644 --- a/crates/storage/provider/provider/src/providers/db/trie.rs +++ b/crates/storage/provider/provider/src/providers/db/trie.rs @@ -1,12 +1,13 @@ -use std::collections::{BTreeMap, HashMap}; +use std::collections::{BTreeMap, BTreeSet, HashMap}; -use katana_db::abstraction::Database; +use katana_db::abstraction::{Database, DbCursor, DbDupSortCursor, DbTx, DbTxRef}; use katana_db::tables; use katana_db::trie::TrieDbMut; use katana_primitives::block::BlockNumber; use katana_primitives::class::{ClassHash, CompiledClassHash}; use katana_primitives::state::StateUpdates; use katana_primitives::{ContractAddress, Felt}; +use katana_provider_api::block::BlockNumberProvider; use katana_provider_api::state::{StateFactoryProvider, StateProvider}; use katana_provider_api::trie::TrieWriter; use katana_provider_api::ProviderError; @@ -108,6 +109,47 @@ impl TrieWriter for DbProvider { Ok(contract_trie_db.root()) })? } + + fn unwind_classes_trie(&self, unwind_to: BlockNumber) -> ProviderResult { + let latest_block_number = self.latest_number()?; + + self.0.update(|tx| { + let mut trie = ClassesTrie::new(TrieDbMut::::new(tx)); + trie.revert_to(unwind_to, latest_block_number); + Ok(trie.root()) + })? + } + + fn unwind_contracts_trie(&self, unwind_to: BlockNumber) -> ProviderResult { + let latest_block_number = self.latest_number()?; + + self.0.update(|tx| { + let mut cursor = tx.cursor_dup::()?; + let iterator = cursor.walk_dup(None, None)?.unwrap(); + + let mut addresses = BTreeSet::new(); + + for entry in iterator { + let (block, change_entry) = entry?; + + if block > unwind_to { + addresses.insert(change_entry.key.contract_address); + } + } + + for addr in addresses { + let trie_db = TrieDbMut::::new(tx); + let mut storage_trie = StoragesTrie::new(trie_db, addr); + storage_trie.revert_to(unwind_to, latest_block_number); + } + + let mut contract_trie_db = + ContractsTrie::new(TrieDbMut::::new(tx)); + contract_trie_db.revert_to(unwind_to, latest_block_number); + + Ok(contract_trie_db.root()) + })? + } } // computes the contract state leaf hash diff --git a/crates/sync/pipeline/src/lib.rs b/crates/sync/pipeline/src/lib.rs index fd1e71cb6..0520694fa 100644 --- a/crates/sync/pipeline/src/lib.rs +++ b/crates/sync/pipeline/src/lib.rs @@ -97,7 +97,9 @@ pub enum Error { #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum PipelineCommand { /// Set the target tip block for the pipeline to sync to. - SetTip(BlockNumber), + Sync(BlockNumber), + /// Set the target tip block for the pipeline to unwind to. + Unwind(BlockNumber), /// Signal the pipeline to stop. Stop, } @@ -121,7 +123,12 @@ impl PipelineHandle { /// /// Panics if the [`Pipeline`] has been dropped. pub fn set_tip(&self, tip: BlockNumber) { - self.tx.send(Some(PipelineCommand::SetTip(tip))).expect("pipeline is no longer running"); + self.tx.send(Some(PipelineCommand::Sync(tip))).expect("pipeline is no longer running"); + } + + pub fn unwind(&self, target: BlockNumber) { + info!(target: "pipeline", %target, "Unwinding pipeline"); + let _ = self.tx.send(Some(PipelineCommand::Unwind(target))); } /// Signals the pipeline to stop gracefully. @@ -168,7 +175,14 @@ pub struct Pipeline

{ stages: Vec>, command_rx: watch::Receiver>, command_tx: watch::Sender>, - tip: Option, + status: PipelineStatus, +} + +#[derive(Debug, Clone)] +enum PipelineStatus { + Idling, + Syncing { tip: BlockNumber, current_target: Option }, + Unwinding { to: BlockNumber, current_target: Option }, } impl

Pipeline

{ @@ -191,7 +205,7 @@ impl

Pipeline

{ command_tx: tx, provider, chunk_size, - tip: None, + status: PipelineStatus::Idling, }; (pipeline, handle) } @@ -248,9 +262,13 @@ impl Pipeline

{ debug!(target: "pipeline", "Received stop command."); break; } - Some(PipelineCommand::SetTip(new_tip)) => { + Some(PipelineCommand::Sync(new_tip)) => { + trace!(target: "pipeline", tip = %new_tip, "Received new tip."); + self.status = PipelineStatus::Syncing { tip: new_tip, current_target: None }; + } + Some(PipelineCommand::Unwind(new_tip)) => { info!(target: "pipeline", tip = %new_tip, "A new tip has been set."); - self.tip = Some(new_tip); + self.status = PipelineStatus::Unwinding { to: new_tip, current_target: None }; } None => {} } @@ -289,7 +307,7 @@ impl Pipeline

{ /// /// Returns an error if any stage execution fails or if the pipeline fails to read the /// checkpoint. - pub async fn run_once(&mut self, to: BlockNumber) -> PipelineResult { + pub async fn execute_once(&mut self, to: BlockNumber) -> PipelineResult { let tip = self.tip.expect("qed; should exist by now"); if self.stages.is_empty() { @@ -353,36 +371,136 @@ impl Pipeline

{ Ok(last_block_processed_list.into_iter().min().unwrap_or(to)) } - /// Run the pipeline loop. - async fn run_loop(&mut self) -> PipelineResult<()> { - let mut current_chunk_tip = self.chunk_size; + pub async fn unwind_once(&mut self, to: BlockNumber) -> PipelineResult { + if self.stages.is_empty() { + return Ok(to); + } - loop { - // Process blocks if we have a tip - if let Some(tip) = self.tip { - let to = current_chunk_tip.min(tip); - let last_block_processed = self.run_once(to).await?; - - if last_block_processed >= tip { - info!(target: "pipeline", %tip, "Finished syncing until tip."); - self.tip = None; - current_chunk_tip = last_block_processed; - } else { - current_chunk_tip = (last_block_processed + self.chunk_size).min(tip); - } + // This is so that lagging stages (ie stage with a checkpoint that is less than the rest of + // the stages) will be executed, in the next cycle of `run_to`, with a `to` value + // whose range from the stages' next checkpoint is equal to the pipeline batch size. + // + // This can actually be done without the allocation, but this makes reasoning about the + // code easier. The majority of the execution time will be spent in `stage.execute` anyway + // so optimizing this doesn't yield significant improvements. + let mut last_block_processed_list: Vec = Vec::with_capacity(self.stages.len()); + + for stage in self.stages.iter_mut() { + let id = stage.id(); + + // Get the checkpoint for the stage, otherwise default to block number 0 + let checkpoint = self.provider.checkpoint(id)?.unwrap_or_default(); + // Skip the stage if the checkpoint is greater than or equal to the target block number + if checkpoint <= to { + info!(target: "pipeline", %id, "Skipping stage."); + last_block_processed_list.push(checkpoint); continue; } - info!(target: "pipeline", "Waiting to receive new tip."); + info!(target: "pipeline", %id, from = %checkpoint, %to, "Unwinding stage."); - // block until a new tip is set - self.command_rx - .wait_for(|c| matches!(c, &Some(PipelineCommand::SetTip(_)))) - .await - .expect("qed; channel closed"); + let input = StageExecutionInput::new(checkpoint, to); + let StageExecutionOutput { last_block_processed } = + stage.execute(&input).await.map_err(|error| Error::StageExecution { id, error })?; + + debug_assert!(last_block_processed <= checkpoint); + + self.provider.set_checkpoint(id, last_block_processed)?; + last_block_processed_list.push(last_block_processed); + + info!(target: "pipeline", %id, from = %checkpoint, %to, "Stage unwinding completed."); + } + + Ok(last_block_processed_list.into_iter().max().unwrap_or(to)) + } + + pub async fn unwind_once(&mut self, to: BlockNumber) -> PipelineResult { + if self.stages.is_empty() { + return Ok(to); + } - yield_now().await; + // This is so that lagging stages (ie stage with a checkpoint that is less than the rest of + // the stages) will be executed, in the next cycle of `run_to`, with a `to` value + // whose range from the stages' next checkpoint is equal to the pipeline batch size. + // + // This can actually be done without the allocation, but this makes reasoning about the + // code easier. The majority of the execution time will be spent in `stage.execute` anyway + // so optimizing this doesn't yield significant improvements. + let mut last_block_processed_list: Vec = Vec::with_capacity(self.stages.len()); + + for stage in self.stages.iter_mut() { + let id = stage.id(); + + // Get the checkpoint for the stage, otherwise default to block number 0 + let checkpoint = self.provider.checkpoint(id)?.unwrap_or_default(); + + // Skip the stage if the checkpoint is greater than or equal to the target block number + if checkpoint <= to { + info!(target: "pipeline", %id, "Skipping stage."); + last_block_processed_list.push(checkpoint); + continue; + } + + info!(target: "pipeline", %id, from = %checkpoint, %to, "Unwinding stage."); + + let input = StageExecutionInput::new(checkpoint, to); + let StageExecutionOutput { last_block_processed } = + stage.execute(&input).await.map_err(|error| Error::StageExecution { id, error })?; + + debug_assert!(last_block_processed <= checkpoint); + + self.provider.set_checkpoint(id, last_block_processed)?; + last_block_processed_list.push(last_block_processed); + + info!(target: "pipeline", %id, from = %checkpoint, %to, "Stage execution completed."); + } + + Ok(last_block_processed_list.into_iter().max().unwrap_or(to)) + } + + /// Run the pipeline loop. + async fn run_loop(&mut self) -> PipelineResult<()> { + loop { + match self.status { + PipelineStatus::Syncing { tip, current_target } => { + let local_to = current_target.unwrap_or(self.chunk_size).min(tip); + let last_block_processed = self.execute_once(local_to).await?; + + if last_block_processed >= tip { + info!(target: "pipeline", %tip, "Finished syncing until tip."); + self.status = PipelineStatus::Idling; + } else { + let new_target = (last_block_processed + self.chunk_size).min(tip); + self.status = + PipelineStatus::Syncing { tip, current_target: Some(new_target) }; + } + } + + PipelineStatus::Unwinding { to, ref mut current_target } => { + let local_to = current_target.unwrap_or(self.chunk_size).max(to); + let last_block_processed = self.unwind_once(local_to).await?; + + if last_block_processed <= to { + info!(target: "pipeline", %to, "Finished unwinding."); + self.status = PipelineStatus::Idling; + } else { + let new_target = (last_block_processed - self.chunk_size).max(to); + self.status = + PipelineStatus::Unwinding { to, current_target: Some(new_target) }; + } + } + + PipelineStatus::Idling => { + // block until a new tip is set + self.command_rx + .wait_for(|c| matches!(c, &Some(PipelineCommand::Sync(_)))) + .await + .expect("qed; channel closed"); + + c2-small-x86-nyc-1 + } + } } } } diff --git a/crates/sync/pipeline/tests/pipeline.rs b/crates/sync/pipeline/tests/pipeline.rs index ec83108e5..8fd9450bd 100644 --- a/crates/sync/pipeline/tests/pipeline.rs +++ b/crates/sync/pipeline/tests/pipeline.rs @@ -146,7 +146,7 @@ async fn run_to_executes_stage_to_target() { pipeline.add_stage(stage); handle.set_tip(5); - let result = pipeline.run_once(5).await.unwrap(); + let result = pipeline.execute_once(5).await.unwrap(); assert_eq!(result, 5); assert_eq!(provider.checkpoint(stage_clone.id()).unwrap(), Some(5)); @@ -170,7 +170,7 @@ async fn run_to_skips_stage_when_checkpoint_equals_target() { pipeline.add_stage(stage); handle.set_tip(5); - let result = pipeline.run_once(5).await.unwrap(); + let result = pipeline.execute_once(5).await.unwrap(); assert_eq!(result, 5); assert_eq!(stage_clone.executions().len(), 0); // Not executed @@ -189,7 +189,7 @@ async fn run_to_skips_stage_when_checkpoint_exceeds_target() { pipeline.add_stage(stage); handle.set_tip(10); - let result = pipeline.run_once(5).await.unwrap(); + let result = pipeline.execute_once(5).await.unwrap(); assert_eq!(result, 10); // Returns the checkpoint assert_eq!(stage_clone.executions().len(), 0); // Not executed @@ -207,7 +207,7 @@ async fn run_to_uses_checkpoint_plus_one_as_from() { provider.set_checkpoint(stage.id(), 3).unwrap(); pipeline.add_stage(stage); handle.set_tip(10); - pipeline.run_once(10).await.unwrap(); + pipeline.execute_once(10).await.unwrap(); let execs = stage_clone.executions(); assert_eq!(execs.len(), 1); @@ -241,7 +241,7 @@ async fn run_to_executes_all_stages_in_order() { ]); handle.set_tip(5); - pipeline.run_once(5).await.unwrap(); + pipeline.execute_once(5).await.unwrap(); // All stages should be executed once because the tip is 5 and the chunk size is 10 assert_eq!(stage1_clone.execution_count(), 1); @@ -279,7 +279,7 @@ async fn run_to_with_mixed_checkpoints() { provider.set_checkpoint(stage2_clone.id(), 3).unwrap(); handle.set_tip(10); - pipeline.run_once(10).await.unwrap(); + pipeline.execute_once(10).await.unwrap(); // Stage1 should be skipped because its checkpoint (10) >= than the tip (10) assert_eq!(stage1_clone.execution_count(), 0); @@ -317,7 +317,7 @@ async fn run_to_returns_minimum_last_block_processed() { ]); handle.set_tip(20); - let result = pipeline.run_once(20).await.unwrap(); + let result = pipeline.execute_once(20).await.unwrap(); // make sure that all the stages were executed once assert_eq!(stage1_clone.execution_count(), 1); @@ -353,7 +353,7 @@ async fn run_to_middle_stage_skip_continues() { provider.set_checkpoint(stage2_clone.id(), 10).unwrap(); handle.set_tip(10); - pipeline.run_once(10).await.unwrap(); + pipeline.execute_once(10).await.unwrap(); // Stage1 and Stage3 should execute assert_eq!(stage1_clone.execution_count(), 1); @@ -526,7 +526,7 @@ async fn stage_execution_error_stops_pipeline() { pipeline.add_stage(stage); handle.set_tip(10); - let result = pipeline.run_once(10).await; + let result = pipeline.execute_once(10).await; assert!(result.is_err()); // Checkpoint should not be set after failure @@ -549,7 +549,7 @@ async fn stage_error_doesnt_affect_subsequent_runs() { pipeline.add_stage(stage2); handle.set_tip(10); - let error = pipeline.run_once(10).await.unwrap_err(); + let error = pipeline.execute_once(10).await.unwrap_err(); let katana_pipeline::Error::StageExecution { id, error } = error else { panic!("Unexpected error type"); @@ -573,7 +573,7 @@ async fn empty_pipeline_returns_target() { // No stages added handle.set_tip(10); - let result = pipeline.run_once(10).await.unwrap(); + let result = pipeline.execute_once(10).await.unwrap(); assert_eq!(result, 10); } @@ -591,7 +591,7 @@ async fn tip_equals_checkpoint_no_execution() { pipeline.add_stage(stage); handle.set_tip(10); - pipeline.run_once(10).await.unwrap(); + pipeline.execute_once(10).await.unwrap(); assert_eq!(executions.lock().unwrap().len(), 0, "Stage1 should not be executed"); } @@ -612,7 +612,7 @@ async fn tip_less_than_checkpoint_skip_all() { pipeline.add_stage(stage); handle.set_tip(20); - let result = pipeline.run_once(10).await.unwrap(); + let result = pipeline.execute_once(10).await.unwrap(); assert_eq!(result, checkpoint); assert_eq!(executions.lock().unwrap().len(), 0, "Stage1 should not be executed"); @@ -660,20 +660,20 @@ async fn stage_checkpoint() { assert_eq!(initial_checkpoint, None); handle.set_tip(5); - pipeline.run_once(5).await.expect("failed to run the pipeline once"); + pipeline.execute_once(5).await.expect("failed to run the pipeline once"); // check that the checkpoint was set let actual_checkpoint = provider.checkpoint("Mock").unwrap(); assert_eq!(actual_checkpoint, Some(5)); handle.set_tip(10); - pipeline.run_once(10).await.expect("failed to run the pipeline once"); + pipeline.execute_once(10).await.expect("failed to run the pipeline once"); // check that the checkpoint was set let actual_checkpoint = provider.checkpoint("Mock").unwrap(); assert_eq!(actual_checkpoint, Some(10)); - pipeline.run_once(10).await.expect("failed to run the pipeline once"); + pipeline.execute_once(10).await.expect("failed to run the pipeline once"); // check that the checkpoint doesn't change let actual_checkpoint = provider.checkpoint("Mock").unwrap(); diff --git a/crates/sync/stage/src/blocks/mod.rs b/crates/sync/stage/src/blocks/mod.rs index e2ef038df..e556ba59b 100644 --- a/crates/sync/stage/src/blocks/mod.rs +++ b/crates/sync/stage/src/blocks/mod.rs @@ -14,8 +14,8 @@ use katana_primitives::state::{StateUpdates, StateUpdatesWithClasses}; use katana_primitives::transaction::{Tx, TxWithHash}; use katana_primitives::Felt; use katana_provider::api::block::{BlockHashProvider, BlockWriter}; -use katana_provider::ProviderError; use katana_provider::api::stage::StageCheckpointProvider; +use katana_provider::ProviderError; use num_traits::ToPrimitive; use starknet::core::types::ResourcePrice; use tracing::{error, info_span, Instrument}; diff --git a/crates/sync/stage/src/lib.rs b/crates/sync/stage/src/lib.rs index 10e413e2f..85d16300b 100644 --- a/crates/sync/stage/src/lib.rs +++ b/crates/sync/stage/src/lib.rs @@ -38,7 +38,7 @@ impl StageExecutionInput { /// /// Panics if `to < from`, as this violates the type's invariant. pub fn new(from: BlockNumber, to: BlockNumber) -> Self { - assert!(to >= from, "Invalid block range: `to` ({to}) must be >= `from` ({from})"); + // assert!(to >= from, "Invalid block range: `to` ({to}) must be >= `from` ({from})"); Self { from, to } } @@ -148,7 +148,9 @@ pub trait Stage: Send + Sync { /// - All data for blocks > `unwind_to` is removed from relevant database tables /// - The stage checkpoint is updated to reflect the unwound state /// - Database invariants are maintained after the unwind operation - fn unwind<'a>(&'a mut self, unwind_to: BlockNumber) -> BoxFuture<'a, StageResult>; + fn unwind<'a>(&'a mut self, unwind_to: BlockNumber) -> BoxFuture<'a, StageResult> { + unimplemented!() + } } #[cfg(test)] @@ -156,6 +158,7 @@ mod tests { use crate::StageExecutionInput; #[tokio::test] + #[ignore] #[should_panic(expected = "Invalid block range")] async fn invalid_range_panics() { // When from > to, the range is invalid and should panic at construction time diff --git a/crates/sync/stage/src/trie.rs b/crates/sync/stage/src/trie.rs index bf6faa953..f0d5e97cd 100644 --- a/crates/sync/stage/src/trie.rs +++ b/crates/sync/stage/src/trie.rs @@ -111,16 +111,7 @@ where } fn unwind<'a>(&'a mut self, unwind_to: BlockNumber) -> BoxFuture<'a, StageResult> { - Box::pin(async move { - warn!( - target: "stage", - id = %self.id(), - unwind_to = %unwind_to, - "StateTrie unwinding not implemented - requires complex trie unwinding logic" - ); - - unimplemented!("StateTrie unwinding requires complex trie unwinding logic") - }) + todo!() } } diff --git a/crates/trie/Cargo.toml b/crates/trie/Cargo.toml index 9a099fbea..c2521bf4c 100644 --- a/crates/trie/Cargo.toml +++ b/crates/trie/Cargo.toml @@ -17,7 +17,7 @@ starknet-types-core.workspace = true thiserror.workspace = true [dependencies.bonsai-trie] -rev = "8b509fd" +rev = "95de4c6" default-features = false features = [ "std" ] git = "https://github.com/dojoengine/bonsai-trie/" diff --git a/crates/trie/src/classes.rs b/crates/trie/src/classes.rs index ff71e3b3b..3fd4010fd 100644 --- a/crates/trie/src/classes.rs +++ b/crates/trie/src/classes.rs @@ -47,6 +47,10 @@ impl ClassesTrie { pub fn multiproof(&mut self, class_hashes: Vec) -> MultiProof { self.trie.multiproof(Self::BONSAI_IDENTIFIER, class_hashes) } + + pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) { + self.trie.revert_to(block.into(), latest_block.into()); + } } impl ClassesTrie diff --git a/crates/trie/src/contracts.rs b/crates/trie/src/contracts.rs index 40b90c40d..90901c761 100644 --- a/crates/trie/src/contracts.rs +++ b/crates/trie/src/contracts.rs @@ -31,6 +31,10 @@ impl ContractsTrie { let keys = addresses.into_iter().map(Felt::from).collect::>(); self.trie.multiproof(Self::BONSAI_IDENTIFIER, keys) } + + pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) { + self.trie.revert_to(block.into(), latest_block.into()); + } } impl ContractsTrie diff --git a/crates/trie/src/lib.rs b/crates/trie/src/lib.rs index 1e49be803..5e55c853a 100644 --- a/crates/trie/src/lib.rs +++ b/crates/trie/src/lib.rs @@ -3,8 +3,8 @@ pub use bonsai::{BitVec, MultiProof, Path, ProofNode}; pub use bonsai_trie::databases::HashMapDb; use bonsai_trie::BonsaiStorage; pub use bonsai_trie::{BonsaiDatabase, BonsaiPersistentDatabase, BonsaiStorageConfig}; -use katana_primitives::class::ClassHash; use katana_primitives::Felt; +use katana_primitives::{block::BlockNumber, class::ClassHash}; use starknet_types_core::hash::{Pedersen, StarkHash}; pub use {bitvec, bonsai_trie as bonsai}; @@ -40,7 +40,7 @@ where pub fn new(db: DB) -> Self { let config = BonsaiStorageConfig { // we have our own implementation of storing trie changes - max_saved_trie_logs: Some(0), + max_saved_trie_logs: None, // in the bonsai-trie crate, this field seems to be only used in rocksdb impl. // i dont understand why would they add a config thats implementation specific ???? // @@ -61,6 +61,10 @@ where let keys = keys.into_iter().map(|key| key.to_bytes_be().as_bits()[5..].to_owned()); self.storage.get_multi_proof(id, keys).expect("failed to get multiproof") } + + pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) { + self.storage.revert_to(block.into(), latest_block.into()).expect("failed to revert trie"); + } } impl BonsaiTrie @@ -169,4 +173,54 @@ mod tests { assert_eq!(result, expected); } + + #[test] + fn test_revert_to() { + use bonsai_trie::databases; + + // the identifier for the trie + const IDENTIFIER: &[u8] = b"test_trie"; + + // Create a BonsaiStorage with in-memory database and trie logs enabled + let bonsai_db = databases::HashMapDb::::default(); + let mut trie = BonsaiTrie::<_, hash::Pedersen>::new(bonsai_db); + + // Insert values at block 0 + trie.insert(IDENTIFIER, Felt::from(1), Felt::from(100)); + trie.insert(IDENTIFIER, Felt::from(2), Felt::from(200)); + trie.commit(0.into()); + let root_at_block_0 = trie.root(IDENTIFIER); + + // Insert more values at block 1 + trie.insert(IDENTIFIER, Felt::from(3), Felt::from(300)); + trie.insert(IDENTIFIER, Felt::from(4), Felt::from(400)); + trie.commit(1.into()); + let root_at_block_1 = trie.root(IDENTIFIER); + + // Roots should be different + assert_ne!(root_at_block_0, root_at_block_1); + + // Insert even more values at block 2 + trie.insert(IDENTIFIER, Felt::from(5), Felt::from(500)); + trie.commit(2.into()); + let root_at_block_2 = trie.root(IDENTIFIER); + + // Roots should be different + assert_ne!(root_at_block_1, root_at_block_2); + assert_ne!(root_at_block_0, root_at_block_2); + + // Revert to block 1 + trie.revert_to(1, 2); + let root_after_revert = trie.root(IDENTIFIER); + + // After revert, root should match block 1 + assert_eq!(root_after_revert, root_at_block_1); + + // Revert to block 0 + trie.revert_to(0, 1); + let root_after_second_revert = trie.root(IDENTIFIER); + + // After revert, root should match block 0 + assert_eq!(root_after_second_revert, root_at_block_0); + } } diff --git a/crates/trie/src/storages.rs b/crates/trie/src/storages.rs index c7f6bbe91..82923f802 100644 --- a/crates/trie/src/storages.rs +++ b/crates/trie/src/storages.rs @@ -25,6 +25,10 @@ impl StoragesTrie { pub fn multiproof(&mut self, storage_keys: Vec) -> MultiProof { self.trie.multiproof(&self.address.to_bytes_be(), storage_keys) } + + pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) { + self.trie.revert_to(block.into(), latest_block.into()); + } } impl StoragesTrie From e3d4201f4622cf329d9c4eb5e5cbc191fe1b791c Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Fri, 17 Oct 2025 12:37:05 -0400 Subject: [PATCH 07/14] wip --- crates/sync/pipeline/src/lib.rs | 56 +++------------------------------ crates/sync/stage/Cargo.toml | 1 + crates/sync/stage/src/trie.rs | 6 +++- 3 files changed, 11 insertions(+), 52 deletions(-) diff --git a/crates/sync/pipeline/src/lib.rs b/crates/sync/pipeline/src/lib.rs index 0520694fa..007486de7 100644 --- a/crates/sync/pipeline/src/lib.rs +++ b/crates/sync/pipeline/src/lib.rs @@ -308,8 +308,6 @@ impl Pipeline

{ /// Returns an error if any stage execution fails or if the pipeline fails to read the /// checkpoint. pub async fn execute_once(&mut self, to: BlockNumber) -> PipelineResult { - let tip = self.tip.expect("qed; should exist by now"); - if self.stages.is_empty() { return Ok(to); } @@ -415,50 +413,6 @@ impl Pipeline

{ Ok(last_block_processed_list.into_iter().max().unwrap_or(to)) } - pub async fn unwind_once(&mut self, to: BlockNumber) -> PipelineResult { - if self.stages.is_empty() { - return Ok(to); - } - - // This is so that lagging stages (ie stage with a checkpoint that is less than the rest of - // the stages) will be executed, in the next cycle of `run_to`, with a `to` value - // whose range from the stages' next checkpoint is equal to the pipeline batch size. - // - // This can actually be done without the allocation, but this makes reasoning about the - // code easier. The majority of the execution time will be spent in `stage.execute` anyway - // so optimizing this doesn't yield significant improvements. - let mut last_block_processed_list: Vec = Vec::with_capacity(self.stages.len()); - - for stage in self.stages.iter_mut() { - let id = stage.id(); - - // Get the checkpoint for the stage, otherwise default to block number 0 - let checkpoint = self.provider.checkpoint(id)?.unwrap_or_default(); - - // Skip the stage if the checkpoint is greater than or equal to the target block number - if checkpoint <= to { - info!(target: "pipeline", %id, "Skipping stage."); - last_block_processed_list.push(checkpoint); - continue; - } - - info!(target: "pipeline", %id, from = %checkpoint, %to, "Unwinding stage."); - - let input = StageExecutionInput::new(checkpoint, to); - let StageExecutionOutput { last_block_processed } = - stage.execute(&input).await.map_err(|error| Error::StageExecution { id, error })?; - - debug_assert!(last_block_processed <= checkpoint); - - self.provider.set_checkpoint(id, last_block_processed)?; - last_block_processed_list.push(last_block_processed); - - info!(target: "pipeline", %id, from = %checkpoint, %to, "Stage execution completed."); - } - - Ok(last_block_processed_list.into_iter().max().unwrap_or(to)) - } - /// Run the pipeline loop. async fn run_loop(&mut self) -> PipelineResult<()> { loop { @@ -471,13 +425,14 @@ impl Pipeline

{ info!(target: "pipeline", %tip, "Finished syncing until tip."); self.status = PipelineStatus::Idling; } else { - let new_target = (last_block_processed + self.chunk_size).min(tip); + let new_target = + last_block_processed.saturating_add(self.chunk_size).min(tip); self.status = PipelineStatus::Syncing { tip, current_target: Some(new_target) }; } } - PipelineStatus::Unwinding { to, ref mut current_target } => { + PipelineStatus::Unwinding { to, current_target } => { let local_to = current_target.unwrap_or(self.chunk_size).max(to); let last_block_processed = self.unwind_once(local_to).await?; @@ -485,7 +440,8 @@ impl Pipeline

{ info!(target: "pipeline", %to, "Finished unwinding."); self.status = PipelineStatus::Idling; } else { - let new_target = (last_block_processed - self.chunk_size).max(to); + let new_target = + last_block_processed.saturating_sub(self.chunk_size).max(to); self.status = PipelineStatus::Unwinding { to, current_target: Some(new_target) }; } @@ -497,8 +453,6 @@ impl Pipeline

{ .wait_for(|c| matches!(c, &Some(PipelineCommand::Sync(_)))) .await .expect("qed; channel closed"); - - c2-small-x86-nyc-1 } } } diff --git a/crates/sync/stage/Cargo.toml b/crates/sync/stage/Cargo.toml index f5caca0ad..393eeb258 100644 --- a/crates/sync/stage/Cargo.toml +++ b/crates/sync/stage/Cargo.toml @@ -7,6 +7,7 @@ version.workspace = true [dependencies] katana-core.workspace = true +katana-trie.workspace = true katana-db.workspace = true katana-executor.workspace = true katana-gateway.workspace = true diff --git a/crates/sync/stage/src/trie.rs b/crates/sync/stage/src/trie.rs index f0d5e97cd..086c5fe6c 100644 --- a/crates/sync/stage/src/trie.rs +++ b/crates/sync/stage/src/trie.rs @@ -111,7 +111,11 @@ where } fn unwind<'a>(&'a mut self, unwind_to: BlockNumber) -> BoxFuture<'a, StageResult> { - todo!() + Box::pin(async move { + self.provider.unwind_classes_trie(unwind_to)?; + self.provider.unwind_contracts_trie(unwind_to)?; + Ok(StageExecutionOutput { last_block_processed: unwind_to }) + }) } } From 4ed12f012931a8cb7ba0859eed8c9804130bb2c2 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Fri, 17 Oct 2025 14:13:21 -0400 Subject: [PATCH 08/14] wip --- Cargo.lock | 1 + bin/katana/Cargo.toml | 1 + bin/katana/src/cli/stage/mod.rs | 6 ++++++ bin/katana/src/cli/stage/unwind.rs | 31 ++++++++++++++++++++++++++++++ crates/cli/src/full.rs | 1 + crates/sync/stage/src/lib.rs | 6 +++++- crates/trie/src/lib.rs | 3 ++- 7 files changed, 47 insertions(+), 2 deletions(-) create mode 100644 bin/katana/src/cli/stage/unwind.rs diff --git a/Cargo.lock b/Cargo.lock index b438df3f5..0fab6224f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5857,6 +5857,7 @@ dependencies = [ "katana-provider", "katana-rpc-client", "katana-rpc-types", + "katana-stage", "katana-utils", "piltover", "proptest", diff --git a/bin/katana/Cargo.toml b/bin/katana/Cargo.toml index 533273616..8ce52256d 100644 --- a/bin/katana/Cargo.toml +++ b/bin/katana/Cargo.toml @@ -16,6 +16,7 @@ katana-primitives.workspace = true katana-rpc-client.workspace = true katana-rpc-types.workspace = true katana-utils.workspace = true +katana-stage.workspace = true anyhow.workspace = true async-trait.workspace = true diff --git a/bin/katana/src/cli/stage/mod.rs b/bin/katana/src/cli/stage/mod.rs index 3a264f105..41294d9fd 100644 --- a/bin/katana/src/cli/stage/mod.rs +++ b/bin/katana/src/cli/stage/mod.rs @@ -1,7 +1,10 @@ use anyhow::Result; use clap::{Args, Subcommand}; +use crate::cli::execute_async; + mod checkpoint; +mod unwind; #[derive(Debug, Args)] #[cfg_attr(test, derive(PartialEq))] @@ -15,12 +18,15 @@ pub struct StageArgs { enum Commands { /// Manage stage checkpoints Checkpoint(checkpoint::CheckpointArgs), + /// Unwind a stage to a previous state + Unwind(unwind::UnwindArgs), } impl StageArgs { pub fn execute(self) -> Result<()> { match self.commands { Commands::Checkpoint(args) => args.execute(), + Commands::Unwind(args) => execute_async(args.execute()), } } } diff --git a/bin/katana/src/cli/stage/unwind.rs b/bin/katana/src/cli/stage/unwind.rs new file mode 100644 index 000000000..b9b4caf4f --- /dev/null +++ b/bin/katana/src/cli/stage/unwind.rs @@ -0,0 +1,31 @@ +use anyhow::Result; +use clap::Args; +use katana_primitives::block::BlockNumber; +use katana_provider::providers::db::DbProvider; +use katana_stage::Stage; + +use crate::cli::db::open_db_rw; + +#[derive(Debug, Args)] +#[cfg_attr(test, derive(PartialEq))] +pub struct UnwindArgs { + /// The stage ID to unwind + #[arg(value_name = "STAGE_ID")] + stage_id: String, + + /// The stage ID to unwind to + #[arg(value_name = "UNWIND_TO")] + unwind_to: BlockNumber, + + /// Path to the database directory. + #[arg(short, long)] + path: String, +} + +impl UnwindArgs { + pub async fn execute(self) -> Result<()> { + let provider = DbProvider::new(open_db_rw(&self.path)?); + katana_stage::StateTrie::new(&provider).unwind(self.unwind_to).await?; + Ok(()) + } +} diff --git a/crates/cli/src/full.rs b/crates/cli/src/full.rs index 623f6b848..1fb1ddcd3 100644 --- a/crates/cli/src/full.rs +++ b/crates/cli/src/full.rs @@ -2,6 +2,7 @@ use std::path::PathBuf; use anyhow::{Context, Result}; pub use clap::Parser; +use katana_node::config::db::DbConfig; use katana_node::config::metrics::MetricsConfig; use katana_node::config::rpc::RpcConfig; use katana_node::full; diff --git a/crates/sync/stage/src/lib.rs b/crates/sync/stage/src/lib.rs index 85d16300b..b0edb15f5 100644 --- a/crates/sync/stage/src/lib.rs +++ b/crates/sync/stage/src/lib.rs @@ -133,6 +133,9 @@ pub trait Stage: Send + Sync { /// specific block. All blocks after the `unwind_to` block should be removed, and the /// resulting database state should be as if the stage had only synced up to `unwind_to`. /// + /// If the `unwind_to` block is larger than the state's checkpoint, this method will be a no-op + /// and should return the checkpoint block number. + /// /// # Arguments /// /// * `unwind_to` - The target block number to unwind to. All blocks after this will be removed. @@ -140,7 +143,8 @@ pub trait Stage: Send + Sync { /// # Returns /// /// A future that resolves to a [`StageResult`] containing [`StageExecutionOutput`] - /// with the last block number after unwinding (should equal `unwind_to`). + /// with the last block number after unwinding or the checkpoint block number (if the stage's + /// checkpoint is smaller than the unwind target). /// /// # Implementation Requirements /// diff --git a/crates/trie/src/lib.rs b/crates/trie/src/lib.rs index 5e55c853a..72d3e9f97 100644 --- a/crates/trie/src/lib.rs +++ b/crates/trie/src/lib.rs @@ -3,8 +3,9 @@ pub use bonsai::{BitVec, MultiProof, Path, ProofNode}; pub use bonsai_trie::databases::HashMapDb; use bonsai_trie::BonsaiStorage; pub use bonsai_trie::{BonsaiDatabase, BonsaiPersistentDatabase, BonsaiStorageConfig}; +use katana_primitives::block::BlockNumber; +use katana_primitives::class::ClassHash; use katana_primitives::Felt; -use katana_primitives::{block::BlockNumber, class::ClassHash}; use starknet_types_core::hash::{Pedersen, StarkHash}; pub use {bitvec, bonsai_trie as bonsai}; From 0510fe91a1c80f2f637787e5c4dad67e12fd9f64 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Fri, 17 Oct 2025 16:08:17 -0400 Subject: [PATCH 09/14] wip --- bin/katana/Cargo.toml | 1 + bin/katana/src/cli/stage/mod.rs | 2 +- crates/sync/pipeline/src/lib.rs | 26 +++++++++++++++++++++----- crates/sync/stage/src/lib.rs | 4 +--- 4 files changed, 24 insertions(+), 9 deletions(-) diff --git a/bin/katana/Cargo.toml b/bin/katana/Cargo.toml index 8ce52256d..589703c3d 100644 --- a/bin/katana/Cargo.toml +++ b/bin/katana/Cargo.toml @@ -17,6 +17,7 @@ katana-rpc-client.workspace = true katana-rpc-types.workspace = true katana-utils.workspace = true katana-stage.workspace = true +katana-provider.workspace = true anyhow.workspace = true async-trait.workspace = true diff --git a/bin/katana/src/cli/stage/mod.rs b/bin/katana/src/cli/stage/mod.rs index 41294d9fd..476a81fcc 100644 --- a/bin/katana/src/cli/stage/mod.rs +++ b/bin/katana/src/cli/stage/mod.rs @@ -26,7 +26,7 @@ impl StageArgs { pub fn execute(self) -> Result<()> { match self.commands { Commands::Checkpoint(args) => args.execute(), - Commands::Unwind(args) => execute_async(args.execute()), + Commands::Unwind(args) => execute_async(args.execute())?, } } } diff --git a/crates/sync/pipeline/src/lib.rs b/crates/sync/pipeline/src/lib.rs index 007486de7..4d1a7bad0 100644 --- a/crates/sync/pipeline/src/lib.rs +++ b/crates/sync/pipeline/src/lib.rs @@ -263,7 +263,7 @@ impl Pipeline

{ break; } Some(PipelineCommand::Sync(new_tip)) => { - trace!(target: "pipeline", tip = %new_tip, "Received new tip."); + debug!(target: "pipeline", tip = %new_tip, "Received new tip."); self.status = PipelineStatus::Syncing { tip: new_tip, current_target: None }; } Some(PipelineCommand::Unwind(new_tip)) => { @@ -389,6 +389,10 @@ impl Pipeline

{ // Get the checkpoint for the stage, otherwise default to block number 0 let checkpoint = self.provider.checkpoint(id)?.unwrap_or_default(); + let span = + info_span!(target: "pipeline", "stage.unwind", stage = %id, current_target = %to); + let enter = span.entered(); + // Skip the stage if the checkpoint is greater than or equal to the target block number if checkpoint <= to { info!(target: "pipeline", %id, "Skipping stage."); @@ -396,14 +400,21 @@ impl Pipeline

{ continue; } + let input = StageExecutionInput::new(checkpoint, to); info!(target: "pipeline", %id, from = %checkpoint, %to, "Unwinding stage."); - let input = StageExecutionInput::new(checkpoint, to); - let StageExecutionOutput { last_block_processed } = - stage.execute(&input).await.map_err(|error| Error::StageExecution { id, error })?; + let span = enter.exit(); + let StageExecutionOutput { last_block_processed } = stage + .execute(&input) + .instrument(span.clone()) + .await + .map_err(|error| Error::StageExecution { id, error })?; debug_assert!(last_block_processed <= checkpoint); + let _enter = span.enter(); + info!(target: "pipeline", from = %checkpoint, %to, "Stage unwinding completed."); + self.provider.set_checkpoint(id, last_block_processed)?; last_block_processed_list.push(last_block_processed); @@ -450,9 +461,14 @@ impl Pipeline

{ PipelineStatus::Idling => { // block until a new tip is set self.command_rx - .wait_for(|c| matches!(c, &Some(PipelineCommand::Sync(_)))) + .wait_for(|c| { + matches!(c, &Some(PipelineCommand::Sync(_))) + || matches!(c, &Some(PipelineCommand::Unwind(_))) + }) .await .expect("qed; channel closed"); + + yield_now().await; } } } diff --git a/crates/sync/stage/src/lib.rs b/crates/sync/stage/src/lib.rs index b0edb15f5..a22d5727c 100644 --- a/crates/sync/stage/src/lib.rs +++ b/crates/sync/stage/src/lib.rs @@ -152,9 +152,7 @@ pub trait Stage: Send + Sync { /// - All data for blocks > `unwind_to` is removed from relevant database tables /// - The stage checkpoint is updated to reflect the unwound state /// - Database invariants are maintained after the unwind operation - fn unwind<'a>(&'a mut self, unwind_to: BlockNumber) -> BoxFuture<'a, StageResult> { - unimplemented!() - } + fn unwind<'a>(&'a mut self, unwind_to: BlockNumber) -> BoxFuture<'a, StageResult>; } #[cfg(test)] From 9e685394da0d8f84f1830e9a2e89fee763076551 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Fri, 17 Oct 2025 16:28:10 -0400 Subject: [PATCH 10/14] wip --- bin/katana/Cargo.toml | 1 - crates/core/src/backend/mod.rs | 8 ++++++++ crates/storage/provider/provider-api/src/trie.rs | 8 ++------ crates/storage/provider/provider/src/lib.rs | 8 ++++++++ crates/storage/provider/provider/src/providers/db/trie.rs | 2 +- .../storage/provider/provider/src/providers/fork/trie.rs | 8 ++++++++ crates/trie/src/classes.rs | 2 +- crates/trie/src/contracts.rs | 2 +- crates/trie/src/storages.rs | 2 +- 9 files changed, 30 insertions(+), 11 deletions(-) diff --git a/bin/katana/Cargo.toml b/bin/katana/Cargo.toml index 589703c3d..8ce52256d 100644 --- a/bin/katana/Cargo.toml +++ b/bin/katana/Cargo.toml @@ -17,7 +17,6 @@ katana-rpc-client.workspace = true katana-rpc-types.workspace = true katana-utils.workspace = true katana-stage.workspace = true -katana-provider.workspace = true anyhow.workspace = true async-trait.workspace = true diff --git a/crates/core/src/backend/mod.rs b/crates/core/src/backend/mod.rs index a26234241..2402096b3 100644 --- a/crates/core/src/backend/mod.rs +++ b/crates/core/src/backend/mod.rs @@ -656,4 +656,12 @@ impl TrieWriter for GenesisTrieWriter { trie.commit(block_number); Ok(trie.root()) } + + fn unwind_classes_trie(&self, _: BlockNumber) -> katana_provider::ProviderResult { + unimplemented!() + } + + fn unwind_contracts_trie(&self, _: BlockNumber) -> katana_provider::ProviderResult { + unimplemented!() + } } diff --git a/crates/storage/provider/provider-api/src/trie.rs b/crates/storage/provider/provider-api/src/trie.rs index d570189c6..a7bb06a06 100644 --- a/crates/storage/provider/provider-api/src/trie.rs +++ b/crates/storage/provider/provider-api/src/trie.rs @@ -21,11 +21,7 @@ pub trait TrieWriter: Send + Sync { state_updates: &StateUpdates, ) -> ProviderResult; - fn unwind_classes_trie(&self, unwind_to: BlockNumber) -> ProviderResult { - unimplemented!() - } + fn unwind_classes_trie(&self, unwind_to: BlockNumber) -> ProviderResult; - fn unwind_contracts_trie(&self, unwind_to: BlockNumber) -> ProviderResult { - unimplemented!() - } + fn unwind_contracts_trie(&self, unwind_to: BlockNumber) -> ProviderResult; } diff --git a/crates/storage/provider/provider/src/lib.rs b/crates/storage/provider/provider/src/lib.rs index c9445f614..877bdc409 100644 --- a/crates/storage/provider/provider/src/lib.rs +++ b/crates/storage/provider/provider/src/lib.rs @@ -375,6 +375,14 @@ where ) -> ProviderResult { self.provider.trie_insert_contract_updates(block_number, state_updates) } + + fn unwind_classes_trie(&self, unwind_to: BlockNumber) -> ProviderResult { + self.provider.unwind_classes_trie(unwind_to) + } + + fn unwind_contracts_trie(&self, unwind_to: BlockNumber) -> ProviderResult { + self.provider.unwind_contracts_trie(unwind_to) + } } impl StageCheckpointProvider for BlockchainProvider diff --git a/crates/storage/provider/provider/src/providers/db/trie.rs b/crates/storage/provider/provider/src/providers/db/trie.rs index 6795232fd..ab8fe3078 100644 --- a/crates/storage/provider/provider/src/providers/db/trie.rs +++ b/crates/storage/provider/provider/src/providers/db/trie.rs @@ -1,6 +1,6 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; -use katana_db::abstraction::{Database, DbCursor, DbDupSortCursor, DbTx, DbTxRef}; +use katana_db::abstraction::{Database, DbDupSortCursor, DbTx}; use katana_db::tables; use katana_db::trie::TrieDbMut; use katana_primitives::block::BlockNumber; diff --git a/crates/storage/provider/provider/src/providers/fork/trie.rs b/crates/storage/provider/provider/src/providers/fork/trie.rs index 743ddcda1..fa081bf03 100644 --- a/crates/storage/provider/provider/src/providers/fork/trie.rs +++ b/crates/storage/provider/provider/src/providers/fork/trie.rs @@ -30,4 +30,12 @@ impl TrieWriter for ForkedProvider { let _ = updates; Ok(Felt::ZERO) } + + fn unwind_classes_trie(&self, unwind_to: BlockNumber) -> ProviderResult { + self.provider.unwind_classes_trie(unwind_to) + } + + fn unwind_contracts_trie(&self, unwind_to: BlockNumber) -> ProviderResult { + self.provider.unwind_contracts_trie(unwind_to) + } } diff --git a/crates/trie/src/classes.rs b/crates/trie/src/classes.rs index 3fd4010fd..bfeb13889 100644 --- a/crates/trie/src/classes.rs +++ b/crates/trie/src/classes.rs @@ -49,7 +49,7 @@ impl ClassesTrie { } pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) { - self.trie.revert_to(block.into(), latest_block.into()); + self.trie.revert_to(block, latest_block); } } diff --git a/crates/trie/src/contracts.rs b/crates/trie/src/contracts.rs index 90901c761..f915ef62a 100644 --- a/crates/trie/src/contracts.rs +++ b/crates/trie/src/contracts.rs @@ -33,7 +33,7 @@ impl ContractsTrie { } pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) { - self.trie.revert_to(block.into(), latest_block.into()); + self.trie.revert_to(block, latest_block); } } diff --git a/crates/trie/src/storages.rs b/crates/trie/src/storages.rs index 82923f802..4b5240271 100644 --- a/crates/trie/src/storages.rs +++ b/crates/trie/src/storages.rs @@ -27,7 +27,7 @@ impl StoragesTrie { } pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) { - self.trie.revert_to(block.into(), latest_block.into()); + self.trie.revert_to(block, latest_block); } } From ad9e3eb20550ec57afb4a560ba86ac49c5cd4acf Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Fri, 17 Oct 2025 21:04:23 -0400 Subject: [PATCH 11/14] wip --- bin/katana/src/cli/stage/unwind.rs | 10 ++++++++-- crates/node/src/full/pool.rs | 13 ++----------- crates/node/src/lib.rs | 2 +- crates/rpc/rpc/src/starknet/blockifier.rs | 2 +- crates/rpc/rpc/src/starknet/mod.rs | 2 +- crates/rpc/rpc/src/starknet/read.rs | 1 - crates/rpc/rpc/src/starknet/write.rs | 8 ++++---- crates/storage/db/src/trie/mod.rs | 19 ++++++++++++++++++- .../provider/src/providers/db/trie.rs | 6 ++++-- crates/sync/stage/src/classes.rs | 2 +- crates/sync/stage/src/trie.rs | 7 ++++--- 11 files changed, 44 insertions(+), 28 deletions(-) diff --git a/bin/katana/src/cli/stage/unwind.rs b/bin/katana/src/cli/stage/unwind.rs index b9b4caf4f..059bd3d10 100644 --- a/bin/katana/src/cli/stage/unwind.rs +++ b/bin/katana/src/cli/stage/unwind.rs @@ -1,7 +1,7 @@ use anyhow::Result; use clap::Args; use katana_primitives::block::BlockNumber; -use katana_provider::providers::db::DbProvider; +use katana_provider::{api::stage::StageCheckpointProvider, providers::db::DbProvider}; use katana_stage::Stage; use crate::cli::db::open_db_rw; @@ -24,8 +24,14 @@ pub struct UnwindArgs { impl UnwindArgs { pub async fn execute(self) -> Result<()> { + use katana_stage::StateTrie; + let provider = DbProvider::new(open_db_rw(&self.path)?); - katana_stage::StateTrie::new(&provider).unwind(self.unwind_to).await?; + let mut stage = StateTrie::new(&provider); + + stage.unwind(self.unwind_to).await?; + provider.set_checkpoint(stage.id(), self.unwind_to)?; + Ok(()) } } diff --git a/crates/node/src/full/pool.rs b/crates/node/src/full/pool.rs index 2eed2e11a..a0f32a854 100644 --- a/crates/node/src/full/pool.rs +++ b/crates/node/src/full/pool.rs @@ -2,18 +2,9 @@ use std::future::Future; use katana_pool::ordering::FiFo; use katana_pool::pool::Pool; -use katana_pool::validation::stateful::TxValidator; -use katana_pool::PoolTransaction; use katana_pool_api::validation::{ValidationOutcome, ValidationResult, Validator}; -use katana_primitives::chain::ChainId; -use katana_primitives::fee::ResourceBoundsMapping; -use katana_primitives::transaction::{ - DeclareTx, DeployAccountTx, ExecutableTxWithHash, InvokeTx, TxHash, -}; -use katana_primitives::{ContractAddress, Felt}; -use katana_rpc_types::{ - BroadcastedDeclareTx, BroadcastedDeployAccountTx, BroadcastedInvokeTx, BroadcastedTx, -}; +use katana_primitives::transaction::ExecutableTxWithHash; +use katana_rpc_types::BroadcastedTx; pub type FullNodePool = Pool>; diff --git a/crates/node/src/lib.rs b/crates/node/src/lib.rs index 3c0856de9..3671a8f81 100644 --- a/crates/node/src/lib.rs +++ b/crates/node/src/lib.rs @@ -31,7 +31,7 @@ use katana_metrics::sys::DiskReporter; use katana_metrics::{Report, Server as MetricsServer}; use katana_pool::ordering::FiFo; use katana_pool::TxPool; -use katana_primitives::env::{FeeTokenAddressses, VersionedConstantsOverrides}; +use katana_primitives::env::VersionedConstantsOverrides; #[cfg(feature = "cartridge")] use katana_rpc::cartridge::CartridgeApi; use katana_rpc::cors::Cors; diff --git a/crates/rpc/rpc/src/starknet/blockifier.rs b/crates/rpc/rpc/src/starknet/blockifier.rs index dcbf6c67b..3e49aff7d 100644 --- a/crates/rpc/rpc/src/starknet/blockifier.rs +++ b/crates/rpc/rpc/src/starknet/blockifier.rs @@ -9,7 +9,7 @@ use katana_executor::implementation::blockifier::utils::{self, block_context_fro use katana_executor::{ExecutionError, ExecutionFlags, ExecutionResult, ResultAndStates}; use katana_primitives::env::{BlockEnv, VersionedConstantsOverrides}; use katana_primitives::transaction::ExecutableTxWithHash; -use katana_primitives::{chain, Felt}; +use katana_primitives::Felt; use katana_provider::api::state::StateProvider; use katana_rpc_api::error::starknet::{ContractErrorData, StarknetApiError}; use katana_rpc_types::{FeeEstimate, FunctionCall}; diff --git a/crates/rpc/rpc/src/starknet/mod.rs b/crates/rpc/rpc/src/starknet/mod.rs index 285d7a73e..e2ffea949 100644 --- a/crates/rpc/rpc/src/starknet/mod.rs +++ b/crates/rpc/rpc/src/starknet/mod.rs @@ -6,7 +6,7 @@ use std::sync::Arc; use katana_chain_spec::ChainSpec; use katana_core::backend::storage::Database; -use katana_pool::{TransactionPool, TxPool}; +use katana_pool::TransactionPool; use katana_primitives::block::{BlockHashOrNumber, BlockIdOrTag, FinalityStatus}; use katana_primitives::class::{ClassHash, CompiledClass}; use katana_primitives::contract::{ContractAddress, Nonce, StorageKey, StorageValue}; diff --git a/crates/rpc/rpc/src/starknet/read.rs b/crates/rpc/rpc/src/starknet/read.rs index 2fa41574e..fedd326d2 100644 --- a/crates/rpc/rpc/src/starknet/read.rs +++ b/crates/rpc/rpc/src/starknet/read.rs @@ -5,7 +5,6 @@ use std::sync::Arc; use anyhow::anyhow; use jsonrpsee::core::{async_trait, RpcResult}; use jsonrpsee::types::ErrorObjectOwned; -use katana_executor::ExecutorFactory; #[cfg(feature = "cartridge")] use katana_genesis::allocation::GenesisAccountAlloc; use katana_pool::TransactionPool; diff --git a/crates/rpc/rpc/src/starknet/write.rs b/crates/rpc/rpc/src/starknet/write.rs index 81cfd0e5a..77b64d9e2 100644 --- a/crates/rpc/rpc/src/starknet/write.rs +++ b/crates/rpc/rpc/src/starknet/write.rs @@ -48,8 +48,8 @@ where .into_inner(this.inner.chain_spec.id()) .map_err(|_| StarknetApiError::InvalidContractClass)?; - let class_hash = tx.class_hash(); - let tx = ExecutableTxWithHash::new(ExecutableTx::Declare(tx)); + let _ = tx.class_hash(); + let _ = ExecutableTxWithHash::new(ExecutableTx::Declare(tx)); // let transaction_hash = this.inner.pool.add_transaction(tx).await?; // Ok(AddDeclareTransactionResponse { transaction_hash, class_hash }) @@ -68,9 +68,9 @@ where } let tx = tx.into_inner(this.inner.chain_spec.id()); - let contract_address = tx.contract_address(); + let _ = tx.contract_address(); - let tx = ExecutableTxWithHash::new(ExecutableTx::DeployAccount(tx)); + let _ = ExecutableTxWithHash::new(ExecutableTx::DeployAccount(tx)); // let transaction_hash = this.inner.pool.add_transaction(tx).await?; // Ok(AddDeployAccountTransactionResponse { transaction_hash, contract_address }) diff --git a/crates/storage/db/src/trie/mod.rs b/crates/storage/db/src/trie/mod.rs index aba812b3a..e4c89ddb1 100644 --- a/crates/storage/db/src/trie/mod.rs +++ b/crates/storage/db/src/trie/mod.rs @@ -163,7 +163,7 @@ where fn create_batch(&self) -> Self::Batch {} fn remove_by_prefix(&mut self, _: &DatabaseKey<'_>) -> Result<(), Self::DatabaseError> { - Ok(()) + unimplemented!() } fn get(&self, key: &DatabaseKey<'_>) -> Result, Self::DatabaseError> { @@ -566,5 +566,22 @@ mod tests { // After revert, root should match block 0 assert_eq!(root_after_second_revert, root_at_block_0); + + // Insert more values at block 1 + trie.insert(felt!("0x3"), felt!("0x300")); + trie.insert(felt!("0x4"), felt!("0x400")); + trie.commit(1); + let root_at_block_1_after_insert = trie.root(); + + // After insertion, root should match block 1 + assert_eq!(root_at_block_1_after_insert, root_at_block_1); + + // Insert even more values at block 2 + trie.insert(felt!("0x5"), felt!("0x500")); + trie.commit(2); + let root_at_block_2_after_insert = trie.root(); + + // After insertion, root should match block 2 + assert_eq!(root_at_block_2_after_insert, root_at_block_2); } } diff --git a/crates/storage/provider/provider/src/providers/db/trie.rs b/crates/storage/provider/provider/src/providers/db/trie.rs index ab8fe3078..6d0c45713 100644 --- a/crates/storage/provider/provider/src/providers/db/trie.rs +++ b/crates/storage/provider/provider/src/providers/db/trie.rs @@ -1,6 +1,6 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; -use katana_db::abstraction::{Database, DbDupSortCursor, DbTx}; +use katana_db::abstraction::{Database, DbCursor, DbDupSortCursor, DbTx}; use katana_db::tables; use katana_db::trie::TrieDbMut; use katana_primitives::block::BlockNumber; @@ -125,7 +125,7 @@ impl TrieWriter for DbProvider { self.0.update(|tx| { let mut cursor = tx.cursor_dup::()?; - let iterator = cursor.walk_dup(None, None)?.unwrap(); + let iterator = cursor.walk(Some(unwind_to))?; let mut addresses = BTreeSet::new(); @@ -137,6 +137,8 @@ impl TrieWriter for DbProvider { } } + dbg!(addresses.len()); + for addr in addresses { let trie_db = TrieDbMut::::new(tx); let mut storage_trie = StoragesTrie::new(trie_db, addr); diff --git a/crates/sync/stage/src/classes.rs b/crates/sync/stage/src/classes.rs index fc1b4d30c..374c8274f 100644 --- a/crates/sync/stage/src/classes.rs +++ b/crates/sync/stage/src/classes.rs @@ -15,7 +15,7 @@ use katana_provider::api::state_update::StateUpdateProvider; use katana_provider::api::ProviderError; use katana_rpc_types::class::ConversionError; use rayon::prelude::*; -use tracing::{debug, error, info, info_span, Instrument}; +use tracing::{debug, error, info_span, Instrument}; use super::{Stage, StageExecutionInput, StageExecutionOutput, StageResult}; use crate::blocks::DatabaseProvider; diff --git a/crates/sync/stage/src/trie.rs b/crates/sync/stage/src/trie.rs index 086c5fe6c..ebdcdf697 100644 --- a/crates/sync/stage/src/trie.rs +++ b/crates/sync/stage/src/trie.rs @@ -1,10 +1,9 @@ use futures::future::BoxFuture; use katana_primitives::block::BlockNumber; use katana_primitives::Felt; -use katana_provider::api::block::HeaderProvider; use katana_provider::api::state_update::StateUpdateProvider; use katana_provider::api::trie::TrieWriter; -use katana_rpc_types::class; +use katana_provider::api::{block::HeaderProvider, state::StateFactoryProvider}; use katana_trie::CommitId; use starknet::macros::short_string; use starknet_types_core::hash::{Poseidon, StarkHash}; @@ -34,7 +33,7 @@ impl

StateTrie

{ impl

Stage for StateTrie

where - P: StateUpdateProvider + TrieWriter + HeaderProvider, + P: StateUpdateProvider + TrieWriter + HeaderProvider + StateFactoryProvider, { fn id(&self) -> &'static str { "StateTrie" @@ -58,6 +57,8 @@ where .state_update(block_number.into())? .ok_or(Error::MissingStateUpdate(block_number))?; + let prev_contract_trie_root = self.provider.latest()?.contracts_root()?; + let computed_contract_trie_root = self.provider.trie_insert_contract_updates(block_number, &state_update)?; From 00c888b5dfa2ba3329e0361aeec52f8f6191dbaf Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Wed, 22 Oct 2025 14:14:29 -0400 Subject: [PATCH 12/14] wip --- bin/katana/src/cli/stage/unwind.rs | 3 ++- crates/storage/provider/provider/src/providers/db/trie.rs | 2 +- crates/sync/stage/src/trie.rs | 3 ++- crates/trie/src/lib.rs | 2 +- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/bin/katana/src/cli/stage/unwind.rs b/bin/katana/src/cli/stage/unwind.rs index 059bd3d10..b7a620f91 100644 --- a/bin/katana/src/cli/stage/unwind.rs +++ b/bin/katana/src/cli/stage/unwind.rs @@ -1,7 +1,8 @@ use anyhow::Result; use clap::Args; use katana_primitives::block::BlockNumber; -use katana_provider::{api::stage::StageCheckpointProvider, providers::db::DbProvider}; +use katana_provider::api::stage::StageCheckpointProvider; +use katana_provider::providers::db::DbProvider; use katana_stage::Stage; use crate::cli::db::open_db_rw; diff --git a/crates/storage/provider/provider/src/providers/db/trie.rs b/crates/storage/provider/provider/src/providers/db/trie.rs index 6d0c45713..9fb718607 100644 --- a/crates/storage/provider/provider/src/providers/db/trie.rs +++ b/crates/storage/provider/provider/src/providers/db/trie.rs @@ -1,6 +1,6 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; -use katana_db::abstraction::{Database, DbCursor, DbDupSortCursor, DbTx}; +use katana_db::abstraction::{Database, DbCursor, DbTx}; use katana_db::tables; use katana_db::trie::TrieDbMut; use katana_primitives::block::BlockNumber; diff --git a/crates/sync/stage/src/trie.rs b/crates/sync/stage/src/trie.rs index ebdcdf697..3b9e4c4d0 100644 --- a/crates/sync/stage/src/trie.rs +++ b/crates/sync/stage/src/trie.rs @@ -1,9 +1,10 @@ use futures::future::BoxFuture; use katana_primitives::block::BlockNumber; use katana_primitives::Felt; +use katana_provider::api::block::HeaderProvider; +use katana_provider::api::state::StateFactoryProvider; use katana_provider::api::state_update::StateUpdateProvider; use katana_provider::api::trie::TrieWriter; -use katana_provider::api::{block::HeaderProvider, state::StateFactoryProvider}; use katana_trie::CommitId; use starknet::macros::short_string; use starknet_types_core::hash::{Poseidon, StarkHash}; diff --git a/crates/trie/src/lib.rs b/crates/trie/src/lib.rs index 72d3e9f97..1839a715f 100644 --- a/crates/trie/src/lib.rs +++ b/crates/trie/src/lib.rs @@ -41,7 +41,7 @@ where pub fn new(db: DB) -> Self { let config = BonsaiStorageConfig { // we have our own implementation of storing trie changes - max_saved_trie_logs: None, + max_saved_trie_logs: Some(0), // in the bonsai-trie crate, this field seems to be only used in rocksdb impl. // i dont understand why would they add a config thats implementation specific ???? // From b643a487b8d0cb45b8929c48d42b5dc3824f96a9 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Wed, 22 Oct 2025 14:33:40 -0400 Subject: [PATCH 13/14] wip --- crates/trie/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/trie/src/lib.rs b/crates/trie/src/lib.rs index 1839a715f..72d3e9f97 100644 --- a/crates/trie/src/lib.rs +++ b/crates/trie/src/lib.rs @@ -41,7 +41,7 @@ where pub fn new(db: DB) -> Self { let config = BonsaiStorageConfig { // we have our own implementation of storing trie changes - max_saved_trie_logs: Some(0), + max_saved_trie_logs: None, // in the bonsai-trie crate, this field seems to be only used in rocksdb impl. // i dont understand why would they add a config thats implementation specific ???? // From 81cb7484b55667e3290ee04b9af7057653a189d1 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Thu, 23 Oct 2025 17:14:50 -0400 Subject: [PATCH 14/14] wip --- crates/storage/provider/provider/src/providers/db/trie.rs | 1 - crates/sync/stage/Cargo.toml | 1 - crates/sync/stage/src/blocks/mod.rs | 2 +- crates/sync/stage/src/classes.rs | 1 - 4 files changed, 1 insertion(+), 4 deletions(-) diff --git a/crates/storage/provider/provider/src/providers/db/trie.rs b/crates/storage/provider/provider/src/providers/db/trie.rs index 9fb718607..3fa5e4ab2 100644 --- a/crates/storage/provider/provider/src/providers/db/trie.rs +++ b/crates/storage/provider/provider/src/providers/db/trie.rs @@ -159,7 +159,6 @@ fn contract_state_leaf_hash( provider: impl StateProvider, address: &ContractAddress, contract_leaf: &ContractLeaf, - block_number: BlockNumber, ) -> Felt { let nonce = contract_leaf.nonce.unwrap_or(provider.nonce(*address).unwrap().unwrap_or_default()); diff --git a/crates/sync/stage/Cargo.toml b/crates/sync/stage/Cargo.toml index 393eeb258..0076a5b14 100644 --- a/crates/sync/stage/Cargo.toml +++ b/crates/sync/stage/Cargo.toml @@ -13,7 +13,6 @@ katana-executor.workspace = true katana-gateway.workspace = true katana-messaging.workspace = true katana-pool.workspace = true -katana-trie.workspace = true katana-primitives.workspace = true katana-provider.workspace = true katana-rpc-types.workspace = true diff --git a/crates/sync/stage/src/blocks/mod.rs b/crates/sync/stage/src/blocks/mod.rs index e556ba59b..6b3714478 100644 --- a/crates/sync/stage/src/blocks/mod.rs +++ b/crates/sync/stage/src/blocks/mod.rs @@ -18,7 +18,7 @@ use katana_provider::api::stage::StageCheckpointProvider; use katana_provider::ProviderError; use num_traits::ToPrimitive; use starknet::core::types::ResourcePrice; -use tracing::{error, info_span, Instrument}; +use tracing::{debug, error, info_span, Instrument}; use crate::{Stage, StageExecutionInput, StageExecutionOutput, StageResult}; diff --git a/crates/sync/stage/src/classes.rs b/crates/sync/stage/src/classes.rs index 374c8274f..e4541a347 100644 --- a/crates/sync/stage/src/classes.rs +++ b/crates/sync/stage/src/classes.rs @@ -199,7 +199,6 @@ where for (key, class) in declared_class_hashes.iter().zip(verified_classes.into_iter()) { self.provider.set_class(key.class_hash, class)?; } - } else { } Ok(StageExecutionOutput { last_block_processed: input.to() })