block downloader separation

Author: kariy

Cargo.lock

@@ -136,9 +136,8 @@ impl NodeArgs {
 
     fn init_logging(&self) -> Result<()> {
         const DEFAULT_LOG_FILTER: &str =
-            "pipeline=debug,info,tasks=debug,executor=trace,forking::backend=trace,\
-                                          blockifier=off,jsonrpsee_server=off,hyper=off,\
-                                          messaging=debug,node=error";
+            "pipeline=debug,info,tasks=debug,executor=trace,forking::backend=trace,blockifier=off,\
+             jsonrpsee_server=off,hyper=off,messaging=debug,node=error";
 
         let filter = if self.development.dev {
             &format!("{DEFAULT_LOG_FILTER},server=debug")

crates/katana/cli/src/args.rs

@@ -12,7 +12,7 @@ use katana_primitives::version::ProtocolVersion;
 use katana_provider::providers::db::DbProvider;
 use katana_provider::providers::fork::ForkedProvider;
 use katana_provider::traits::block::{BlockProvider, BlockWriter};
-use katana_provider::traits::contract::ContractClassWriter;
+use katana_provider::traits::contract::{ContractClassWriter, ContractClassWriterExt};
 use katana_provider::traits::env::BlockEnvProvider;
 use katana_provider::traits::stage::StageCheckpointProvider;
 use katana_provider::traits::state::{StateFactoryProvider, StateRootProvider, StateWriter};
@@ -43,6 +43,7 @@ pub trait Database:
     + StateRootProvider
     + StateWriter
     + ContractClassWriter
+    + ContractClassWriterExt
     + StateFactoryProvider
     + BlockEnvProvider
     + ClassTrieWriter
@@ -67,6 +68,7 @@ impl<T> Database for T where
         + StateRootProvider
         + StateWriter
         + ContractClassWriter
+        + ContractClassWriterExt
         + StateFactoryProvider
         + BlockEnvProvider
         + ClassTrieWriter

crates/katana/core/src/backend/storage.rs

@@ -11,14 +11,16 @@ katana-executor.workspace = true
 katana-pool.workspace = true
 katana-primitives.workspace = true
 katana-provider = { workspace = true, features = [ "test-utils" ] }
+katana-rpc-types.workspace = true
 katana-tasks.workspace = true
 
 anyhow.workspace = true
 async-trait.workspace = true
+backon = { version = "1.3.0", features = [ "tokio-sleep" ] }
 futures.workspace = true
+parking_lot = "0.12.3"
+serde_json = "1.0.133"
 starknet.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
 tracing.workspace = true
-serde_json = "1.0.133"
-parking_lot = "0.12.3"

crates/katana/pipeline/Cargo.toml

@@ -1,10 +1,12 @@
-use std::time::Duration;
+use std::sync::Arc;
 
+use backon::{ExponentialBuilder, Retryable};
 use katana_primitives::block::{BlockNumber, SealedBlockWithStatus};
 use katana_primitives::state::{StateUpdates, StateUpdatesWithClasses};
 use katana_provider::traits::block::BlockWriter;
 use starknet::providers::sequencer::models::{BlockId, StateUpdateWithBlock};
 use starknet::providers::{ProviderError, SequencerGatewayProvider};
+use tracing::warn;
 
 use super::{Stage, StageExecutionInput, StageResult};
 
@@ -17,20 +19,12 @@ pub enum Error {
 #[derive(Debug)]
 pub struct Blocks<P> {
     provider: P,
-    feeder_gateway: SequencerGatewayProvider,
+    downloader: Downloader,
 }
 
 impl<P> Blocks<P> {
     pub fn new(provider: P, feeder_gateway: SequencerGatewayProvider) -> Self {
-        Self { provider, feeder_gateway }
-    }
-}
-
-impl<P: BlockWriter> Blocks<P> {
-    #[allow(deprecated)]
-    async fn fetch_block(&self, block: BlockNumber) -> Result<StateUpdateWithBlock, Error> {
-        let res = self.feeder_gateway.get_state_update_with_block(BlockId::Number(block)).await?;
-        Ok(res)
+        Self { provider, downloader: Downloader::new(feeder_gateway) }
     }
 }
 
@@ -41,10 +35,11 @@ impl<P: BlockWriter> Stage for Blocks<P> {
     }
 
     async fn execute(&mut self, input: &StageExecutionInput) -> StageResult {
-        let mut current_block = input.from;
+        // Download all blocks concurrently
+        let blocks = self.downloader.fetch_blocks_range(input.from, input.to, 10).await?;
 
-        loop {
-            let data = self.fetch_block(current_block).await?;
+        // Then process them sequentially
+        for data in blocks {
             let StateUpdateWithBlock { state_update, block: fgw_block } = data;
 
             let block = SealedBlockWithStatus::from(fgw_block);
@@ -57,17 +52,74 @@ impl<P: BlockWriter> Stage for Blocks<P> {
                 Vec::new(),
                 Vec::new(),
             );
+        }
 
-            tokio::time::sleep(Duration::from_secs(1)).await;
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone)]
+struct Downloader {
+    client: Arc<SequencerGatewayProvider>,
+}
 
-            if current_block == input.to {
-                break;
-            } else {
-                current_block += 1;
+impl Downloader {
+    fn new(client: SequencerGatewayProvider) -> Self {
+        Self { client: Arc::new(client) }
+    }
+
+    /// Fetch blocks in the range [from, to] in batches of `batch_size`.
+    async fn fetch_blocks_range(
+        &self,
+        from: BlockNumber,
+        to: BlockNumber,
+        batch_size: usize,
+    ) -> Result<Vec<StateUpdateWithBlock>, Error> {
+        let mut all_results = Vec::with_capacity(to.saturating_sub(from) as usize);
+
+        for batch_start in (from..=to).step_by(batch_size) {
+            let batch_end = (batch_start + batch_size as u64 - 1).min(to);
+
+            // fetch in batches and wait on them before proceeding to the next batch
+            let mut futures = Vec::new();
+            for block_num in batch_start..=batch_end {
+                futures.push(self.fetch_block_with_retry(block_num));
             }
+
+            let batch_results = futures::future::join_all(futures).await;
+            all_results.extend(batch_results);
         }
 
-        Ok(())
+        all_results.into_iter().collect()
+    }
+
+    /// Fetch a single block with the given block number with retry mechanism.
+    async fn fetch_block_with_retry(
+        &self,
+        block: BlockNumber,
+    ) -> Result<StateUpdateWithBlock, Error> {
+        let request = || async move {
+            #[allow(deprecated)]
+            self.clone().fetch_block(block).await
+        };
+
+        // Retry only when being rate limited
+        let result = request
+            .retry(ExponentialBuilder::default())
+            .when(|e| matches!(e, Error::Gateway(ProviderError::RateLimited)))
+            .notify(|error, _| {
+                warn!(target: "pipeline", %block, %error, "Retrying block download.");
+            })
+            .await?;
+
+        Ok(result)
+    }
+
+    /// Fetch a single block with the given block number.
+    async fn fetch_block(&self, block: BlockNumber) -> Result<StateUpdateWithBlock, Error> {
+        #[allow(deprecated)]
+        let res = self.client.get_state_update_with_block(BlockId::Number(block)).await?;
+        Ok(res)
     }
 }
 

crates/katana/pipeline/src/stage/blocks.rs

@@ -2,13 +2,16 @@ use std::time::Duration;
 
 use anyhow::Result;
 use katana_primitives::block::BlockNumber;
-use katana_primitives::class::{CasmContractClass, ClassHash, CompiledClass, ContractClass};
+use katana_primitives::class::{
+    CasmContractClass, ClassHash, CompiledClass, ContractClass, SierraContractClass,
+};
 use katana_primitives::conversion::rpc::{legacy_rpc_to_class, StarknetRsLegacyContractClass};
-use katana_provider::traits::contract::ContractClassWriter;
+use katana_provider::traits::contract::{ContractClassWriter, ContractClassWriterExt};
 use katana_provider::traits::state_update::StateUpdateProvider;
+use katana_rpc_types::class::RpcSierraContractClass;
 use starknet::providers::sequencer::models::{BlockId, DeployedClass};
 use starknet::providers::{ProviderError, SequencerGatewayProvider};
-use tracing::debug;
+use tracing::info;
 
 use super::{Stage, StageExecutionInput, StageResult};
 
@@ -35,44 +38,31 @@ where
     P: StateUpdateProvider + ContractClassWriter,
 {
     #[allow(deprecated)]
-    async fn get_class(
-        &self,
-        hash: ClassHash,
-        block: BlockNumber,
-    ) -> Result<(ContractClass, Option<CompiledClass>), Error> {
-        let block_id = BlockId::Number(block);
+    async fn get_class(&self, hash: ClassHash, block: BlockNumber) -> Result<ContractClass, Error> {
+        let class = self.feeder_gateway.get_class_by_hash(hash, BlockId::Number(block)).await?;
 
-        let (class, casm) = tokio::join!(
-            self.feeder_gateway.get_class_by_hash(hash, block_id),
-            self.feeder_gateway.get_compiled_class_by_class_hash(hash, block_id)
-        );
-
-        let (class, casm) = (class?, casm?);
-
-        let (class, casm) = match class {
+        let class = match class {
             DeployedClass::LegacyClass(legacy) => {
                 let class = to_inner_legacy_class(legacy).unwrap();
-                (class, None)
+                class
             }
 
+            // TODO: implement our own fgw client using our own types for easier conversion
             DeployedClass::SierraClass(sierra) => {
-                // TODO: change this shyte
-                let value = serde_json::to_value(casm).unwrap();
-                let casm = serde_json::from_value::<CasmContractClass>(value).unwrap();
-                let compiled = CompiledClass::Class(casm);
-
-                (ContractClass::Class(sierra), Some(compiled))
+                let rpc_class = RpcSierraContractClass::try_from(sierra).unwrap();
+                let class = SierraContractClass::try_from(rpc_class).unwrap();
+                ContractClass::Class(class)
             }
         };
 
-        Ok((class, casm))
+        Ok(class)
     }
 }
 
 #[async_trait::async_trait]
 impl<P> Stage for Classes<P>
 where
-    P: StateUpdateProvider + ContractClassWriter,
+    P: StateUpdateProvider + ContractClassWriter + ContractClassWriterExt,
 {
     fn id(&self) -> &'static str {
         "Classes"
@@ -84,16 +74,12 @@ where
             let class_hashes = self.provider.declared_classes(i.into())?.unwrap();
 
             // TODO: do this in parallel
-            for class_hash in class_hashes.keys() {
-                debug!(target: "pipeline", "Fetching class artifacts for class hash {class_hash:#x}");
+            for hash in class_hashes.keys() {
+                info!(target: "pipeline", class_hash = format!("{hash:#x}"), "Fetching class artifacts.");
 
                 // 1. fetch sierra and casm class from fgw
-                let (class, compiled) = self.get_class(*class_hash, i).await?;
-
-                self.provider.set_class(*class_hash, class)?;
-                if let Some(casm) = compiled {
-                    self.provider.set_compiled_class(*class_hash, casm)?;
-                }
+                let class = self.get_class(*hash, i).await?;
+                self.provider.set_class(*hash, class)?;
 
                 tokio::time::sleep(Duration::from_secs(1)).await;
             }

crates/katana/pipeline/src/stage/classes.rs

@@ -163,8 +163,8 @@ fn compute_legacy_class_hash(class: &LegacyContractClass) -> Result<Felt, Comput
 #[cfg(test)]
 mod tests {
 
-    use starknet::core::types::contract::SierraClass as StarknetRsSierraContractClass;
     use starknet::core::types::contract::legacy::LegacyContractClass as StarknetRsLegacyContractClass;
+    use starknet::core::types::contract::SierraClass as StarknetRsSierraContractClass;
 
     use super::{ContractClass, LegacyContractClass, SierraContractClass};
 

crates/katana/primitives/src/class.rs

@@ -1,10 +1,10 @@
 #[cfg(feature = "postcard")]
 pub mod postcard;
 
-use katana_primitives::Felt;
 use katana_primitives::block::FinalityStatus;
 use katana_primitives::class::ContractClass;
 use katana_primitives::contract::ContractAddress;
+use katana_primitives::Felt;
 
 use crate::error::CodecError;