update

kariy · kariy · commit 2db8472bd94a · 2024-11-26T21:47:50.000+08:00
diff --git a/crates/katana/cli/src/args.rs b/crates/katana/cli/src/args.rs
@@ -22,7 +22,6 @@ use serde::{Deserialize, Serialize};
 use tracing::{info, Subscriber};
 use tracing_log::LogTracer;
 use tracing_subscriber::{fmt, EnvFilter};
-use url::Url;
 
 use crate::file::NodeArgsConfig;
 use crate::options::*;
@@ -136,8 +135,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,stage=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")
diff --git a/crates/katana/node/src/lib.rs b/crates/katana/node/src/lib.rs
@@ -33,7 +33,7 @@ use katana_db::mdbx::DbEnv;
 use katana_executor::implementation::blockifier::BlockifierFactory;
 use katana_executor::{ExecutionFlags, ExecutorFactory};
 use katana_pipeline::stage::{Blocks, Classes};
-use katana_pipeline::{stage, Pipeline, PipelineHandle};
+use katana_pipeline::{Pipeline, PipelineHandle};
 use katana_pool::ordering::FiFo;
 use katana_pool::validation::stateful::TxValidator;
 use katana_pool::TxPool;
@@ -134,12 +134,12 @@ impl Node {
 
         let provider = self.backend.blockchain.provider().clone();
         let fgw = SequencerGatewayProvider::starknet_alpha_sepolia();
-        let (mut pipeline, handle) = Pipeline::new(provider.clone(), 10);
+        let (mut pipeline, handle) = Pipeline::new(provider.clone(), 64);
 
-        pipeline.add_stage(Blocks::new(provider.clone(), fgw.clone()));
-        pipeline.add_stage(Classes::new(provider, fgw.clone()));
+        pipeline.add_stage(Blocks::new(provider.clone(), fgw.clone(), 3));
+        pipeline.add_stage(Classes::new(provider, fgw.clone(), 3));
 
-        handle.set_tip(100);
+        handle.set_tip(1000);
 
         self.task_manager
             .task_spawner()
diff --git a/crates/katana/pipeline/src/lib.rs b/crates/katana/pipeline/src/lib.rs
@@ -132,8 +132,9 @@ impl<P: StageCheckpointProvider> Pipeline<P> {
             let input = StageExecutionInput { from: checkpoint + 1, to };
             stage.execute(&input).await?;
             self.provider.set_checkpoint(id, to)?;
-        }
 
+            info!(target: "pipeline", %id, from = %checkpoint, %to, "Stage execution completed.");
+        }
         Ok(())
     }
 }
@@ -193,7 +194,7 @@ mod tests {
     async fn stage_checkpoint() {
         let provider = test_provider();
 
-        let (mut pipeline, handle) = Pipeline::new(&provider, 10);
+        let (mut pipeline, _handle) = Pipeline::new(&provider, 10);
         pipeline.add_stage(MockStage);
 
         // check that the checkpoint was set
diff --git a/crates/katana/pipeline/src/stage/blocks.rs b/crates/katana/pipeline/src/stage/blocks.rs
@@ -1,12 +1,13 @@
 use std::sync::Arc;
+use std::time::Duration;
 
 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 tracing::{debug, warn};
 
 use super::{Stage, StageExecutionInput, StageResult};
 
@@ -23,8 +24,13 @@ pub struct Blocks<P> {
 }
 
 impl<P> Blocks<P> {
-    pub fn new(provider: P, feeder_gateway: SequencerGatewayProvider) -> Self {
-        Self { provider, downloader: Downloader::new(feeder_gateway) }
+    pub fn new(
+        provider: P,
+        feeder_gateway: SequencerGatewayProvider,
+        download_batch_size: usize,
+    ) -> Self {
+        let downloader = Downloader::new(feeder_gateway, download_batch_size);
+        Self { provider, downloader }
     }
 }
 
@@ -36,22 +42,25 @@ impl<P: BlockWriter> Stage for Blocks<P> {
 
     async fn execute(&mut self, input: &StageExecutionInput) -> StageResult {
         // Download all blocks concurrently
-        let blocks = self.downloader.fetch_blocks_range(input.from, input.to, 10).await?;
-
-        // Then process them sequentially
-        for data in blocks {
-            let StateUpdateWithBlock { state_update, block: fgw_block } = data;
-
-            let block = SealedBlockWithStatus::from(fgw_block);
-            let su = StateUpdates::from(state_update);
-            let su = StateUpdatesWithClasses { state_updates: su, ..Default::default() };
-
-            let _ = self.provider.insert_block_with_states_and_receipts(
-                block,
-                su,
-                Vec::new(),
-                Vec::new(),
-            );
+        let blocks = self.downloader.download_blocks(input.from, input.to).await?;
+
+        if !blocks.is_empty() {
+            debug!(target: "stage", id = %self.id(), total = %blocks.len(), "Storing blocks to storage.");
+            // Store blocks to storage
+            for block in blocks {
+                let StateUpdateWithBlock { state_update, block: fgw_block } = block;
+
+                let block = SealedBlockWithStatus::from(fgw_block);
+                let su = StateUpdates::from(state_update);
+                let su = StateUpdatesWithClasses { state_updates: su, ..Default::default() };
+
+                let _ = self.provider.insert_block_with_states_and_receipts(
+                    block,
+                    su,
+                    Vec::new(),
+                    Vec::new(),
+                );
+            }
         }
 
         Ok(())
@@ -60,61 +69,69 @@ impl<P: BlockWriter> Stage for Blocks<P> {
 
 #[derive(Debug, Clone)]
 struct Downloader {
+    batch_size: usize,
     client: Arc<SequencerGatewayProvider>,
 }
 
 impl Downloader {
-    fn new(client: SequencerGatewayProvider) -> Self {
-        Self { client: Arc::new(client) }
+    fn new(client: SequencerGatewayProvider, batch_size: usize) -> Self {
+        Self { client: Arc::new(client), batch_size }
     }
 
     /// Fetch blocks in the range [from, to] in batches of `batch_size`.
-    async fn fetch_blocks_range(
+    async fn download_blocks(
         &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));
-            }
+        debug!(target: "pipeline", %from, %to, "Downloading blocks.");
+        let mut blocks = Vec::with_capacity(to.saturating_sub(from) as usize);
 
-            let batch_results = futures::future::join_all(futures).await;
-            all_results.extend(batch_results);
+        for batch_start in (from..=to).step_by(self.batch_size) {
+            let batch_end = (batch_start + self.batch_size as u64 - 1).min(to);
+            let batch = self.fetch_blocks_with_retry(batch_start, batch_end).await?;
+            blocks.extend(batch);
         }
 
-        all_results.into_iter().collect()
+        Ok(blocks)
     }
 
-    /// Fetch a single block with the given block number with retry mechanism.
-    async fn fetch_block_with_retry(
+    /// Fetch blocks with the given block number with retry mechanism at a batch level.
+    async fn fetch_blocks_with_retry(
         &self,
-        block: BlockNumber,
-    ) -> Result<StateUpdateWithBlock, Error> {
-        let request = || async move {
-            #[allow(deprecated)]
-            self.clone().fetch_block(block).await
-        };
+        from: BlockNumber,
+        to: BlockNumber,
+    ) -> Result<Vec<StateUpdateWithBlock>, Error> {
+        let request = || async move { self.clone().fetch_blocks(from, to).await };
 
         // Retry only when being rate limited
+        let backoff = ExponentialBuilder::default().with_min_delay(Duration::from_secs(9));
         let result = request
-            .retry(ExponentialBuilder::default())
-            .when(|e| matches!(e, Error::Gateway(ProviderError::RateLimited)))
+            .retry(backoff)
             .notify(|error, _| {
-                warn!(target: "pipeline", %block, %error, "Retrying block download.");
+                warn!(target: "pipeline", %from, %to, %error, "Retrying block download.");
             })
             .await?;
 
         Ok(result)
     }
 
+    async fn fetch_blocks(
+        &self,
+        from: BlockNumber,
+        to: BlockNumber,
+    ) -> Result<Vec<StateUpdateWithBlock>, Error> {
+        let total = to.saturating_sub(from) as usize;
+        let mut requests = Vec::with_capacity(total);
+
+        for i in from..=to {
+            requests.push(self.fetch_block(i));
+        }
+
+        let results = futures::future::join_all(requests).await;
+        results.into_iter().collect()
+    }
+
     /// Fetch a single block with the given block number.
     async fn fetch_block(&self, block: BlockNumber) -> Result<StateUpdateWithBlock, Error> {
         #[allow(deprecated)]
@@ -140,7 +157,7 @@ mod tests {
         let provider = test_provider();
         let feeder_gateway = SequencerGatewayProvider::starknet_alpha_sepolia();
 
-        let mut stage = Blocks::new(&provider, feeder_gateway);
+        let mut stage = Blocks::new(&provider, feeder_gateway, 10);
 
         let input = StageExecutionInput { from: from_block, to: to_block };
         let _ = stage.execute(&input).await.expect("failed to execute stage");
diff --git a/crates/katana/pipeline/src/stage/classes.rs b/crates/katana/pipeline/src/stage/classes.rs
@@ -1,4 +1,5 @@
 use std::sync::Arc;
+use std::time::Duration;
 
 use anyhow::Result;
 use backon::{ExponentialBuilder, Retryable};
@@ -10,7 +11,7 @@ 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::warn;
+use tracing::{debug, warn};
 
 use super::{Stage, StageExecutionInput, StageResult};
 
@@ -27,8 +28,13 @@ pub struct Classes<P> {
 }
 
 impl<P> Classes<P> {
-    pub fn new(provider: P, feeder_gateway: SequencerGatewayProvider) -> Self {
-        Self { provider, downloader: Downloader::new(feeder_gateway) }
+    pub fn new(
+        provider: P,
+        feeder_gateway: SequencerGatewayProvider,
+        download_batch_size: usize,
+    ) -> Self {
+        let downloader = Downloader::new(feeder_gateway, download_batch_size);
+        Self { provider, downloader }
     }
 }
 
@@ -43,12 +49,18 @@ where
 
     async fn execute(&mut self, input: &StageExecutionInput) -> StageResult {
         for i in input.from..=input.to {
+            // get the classes declared at block `i`
             let class_hashes = self.provider.declared_classes(i.into())?.unwrap();
             let class_hashes = class_hashes.keys().map(|hash| *hash).collect::<Vec<_>>();
 
-            let classes = self.downloader.fetch_classes(&class_hashes, i).await?;
-            for (hash, class) in classes {
-                self.provider.set_class(hash, class)?;
+            // fetch the classes artifacts
+            let classes = self.downloader.download_classes(&class_hashes, i).await?;
+
+            if !classes.is_empty() {
+                debug!(target: "stage", id = %self.id(), total = %classes.len(), "Storing classes to storage.");
+                for (hash, class) in class_hashes.iter().zip(classes) {
+                    self.provider.set_class(*hash, class)?;
+                }
             }
         }
 
@@ -58,53 +70,63 @@ where
 
 #[derive(Debug, Clone)]
 struct Downloader {
+    batch_size: usize,
     client: Arc<SequencerGatewayProvider>,
 }
 
 impl Downloader {
-    fn new(client: SequencerGatewayProvider) -> Self {
-        Self { client: Arc::new(client) }
+    fn new(client: SequencerGatewayProvider, batch_size: usize) -> Self {
+        Self { client: Arc::new(client), batch_size }
     }
 
-    async fn fetch_classes(
+    async fn download_classes(
         &self,
-        classes: &[ClassHash],
+        hashes: &[ClassHash],
         block: BlockNumber,
-    ) -> Result<Vec<(ClassHash, ContractClass)>, Error> {
-        let mut all_results = Vec::with_capacity(classes.len());
-
-        for hash in classes {
-            let mut futures = Vec::new();
-
-            futures.push(self.fetch_class_with_retry(*hash, block));
-            let batch_results = futures::future::join_all(futures).await;
+    ) -> Result<Vec<ContractClass>, Error> {
+        debug!(total = %hashes.len(), %block, "Downloading classes.");
+        let mut classes = Vec::with_capacity(hashes.len());
 
-            all_results.extend(batch_results);
+        for chunk in hashes.chunks(self.batch_size) {
+            let batch = self.fetch_classes_with_retry(chunk, block).await?;
+            classes.extend(batch);
         }
 
-        all_results.into_iter().collect()
+        Ok(classes)
     }
 
-    async fn fetch_class_with_retry(
+    async fn fetch_classes_with_retry(
         &self,
-        hash: ClassHash,
+        classes: &[ClassHash],
         block: BlockNumber,
-    ) -> Result<(ClassHash, ContractClass), Error> {
-        let request = || async move {
-            #[allow(deprecated)]
-            self.clone().fetch_class(hash, block).await
-        };
+    ) -> Result<Vec<ContractClass>, Error> {
+        let request = || async move { self.clone().fetch_classes(classes, block).await };
 
         // Retry only when being rate limited
+        let backoff = ExponentialBuilder::default().with_min_delay(Duration::from_secs(3));
         let result = request
-            .retry(ExponentialBuilder::default())
-            .when(|e| matches!(e, Error::Gateway(ProviderError::RateLimited)))
+            .retry(backoff)
             .notify(|error, _| {
-                warn!(target: "pipeline", hash = format!("{hash:#x}"), %block, %error, "Retrying class download.");
+                warn!(target: "pipeline", %error, "Retrying class download.");
             })
             .await?;
 
-        Ok((hash, result))
+        Ok(result)
+    }
+
+    async fn fetch_classes(
+        &self,
+        classes: &[ClassHash],
+        block: BlockNumber,
+    ) -> Result<Vec<ContractClass>, Error> {
+        let mut requests = Vec::with_capacity(classes.len());
+
+        for class in classes {
+            requests.push(self.fetch_class(*class, block));
+        }
+
+        let results = futures::future::join_all(requests).await;
+        results.into_iter().collect()
     }
 
     async fn fetch_class(
diff --git a/crates/katana/pipeline/tests/sync.rs b/crates/katana/pipeline/tests/sync.rs
@@ -16,8 +16,8 @@ async fn fgw_sync() {
     // build stages
 
     let fgw = SequencerGatewayProvider::starknet_alpha_sepolia();
-    let blocks = stage::Blocks::new(db_provider.clone(), fgw.clone());
-    let classes = stage::Classes::new(db_provider.clone(), fgw);
+    let blocks = stage::Blocks::new(db_provider.clone(), fgw.clone(), 10);
+    let classes = stage::Classes::new(db_provider.clone(), fgw, 10);
 
     let (mut pipeline, handle) = Pipeline::new(db_provider.clone(), chunk_size);
     pipeline.add_stage(blocks);
