refactor(sharding): schedule shards by id via pool

Author: kariy

Cargo.lock

@@ -20,7 +20,7 @@ katana-metrics.workspace = true
 katana-pool.workspace = true
 katana-primitives.workspace = true
 katana-provider.workspace = true
-katana-rpc-server.workspace = true
+katana-rpc-server = { workspace = true, features = [ "cartridge" ] }
 katana-rpc-api.workspace = true
 katana-rpc-client.workspace = true
 katana-rpc-types.workspace = true

crates/node/sharding/Cargo.toml

@@ -11,7 +11,6 @@ use katana_executor::blockifier::cache::ClassCache;
 use katana_executor::blockifier::BlockifierFactory;
 use katana_executor::{ExecutionFlags, ExecutorFactory};
 use katana_gas_price_oracle::GasPriceOracle;
-use katana_pool::TxPool;
 use katana_primitives::env::VersionedConstantsOverrides;
 use katana_primitives::{ContractAddress, Felt};
 use katana_provider::DbProviderFactory;
@@ -20,6 +19,7 @@ use katana_rpc_server::shard::{ShardProvider, ShardRpc};
 use katana_rpc_server::starknet::{StarknetApi, StarknetApiConfig};
 use katana_rpc_server::{RpcServer, RpcServerHandle};
 use katana_sharding::manager::{LazyShardManager, ShardManager};
+use katana_sharding::pool::ShardPool;
 use katana_sharding::runtime::{Runtime, RuntimeHandle};
 use katana_sharding::shard::NoPendingBlockProvider;
 use katana_tasks::TaskManager;
@@ -103,13 +103,13 @@ impl Node {
             gas_oracle,
             starknet_api_config,
             task_spawner.clone(),
+            handle.scheduler().clone(),
             config.base_chain_url.clone(),
         ));
 
         // --- Build RPC server with shard API
         let provider = NodeShardProvider {
             manager: manager.clone(),
-            handle: handle.clone(),
             chain_spec: config.chain.clone(),
         };
 
@@ -184,12 +184,11 @@ impl LaunchedShardNode {
 /// manager and scheduler to the generic `ShardRpc` RPC handler.
 struct NodeShardProvider {
     manager: Arc<dyn ShardManager>,
-    handle: RuntimeHandle,
     chain_spec: Arc<ChainSpec>,
 }
 
 impl ShardProvider for NodeShardProvider {
-    type Api = StarknetApi<TxPool, NoPendingBlockProvider, DbProviderFactory>;
+    type Api = StarknetApi<ShardPool, NoPendingBlockProvider, DbProviderFactory>;
 
     fn starknet_api(&self, shard_id: ContractAddress) -> Result<Self::Api, ErrorObjectOwned> {
         let shard = self.manager.get(shard_id).map_err(|e| {
@@ -202,12 +201,6 @@ impl ShardProvider for NodeShardProvider {
         Ok(shard.starknet_api.clone())
     }
 
-    fn schedule(&self, shard_id: ContractAddress) {
-        if let Ok(shard) = self.manager.get(shard_id) {
-            self.handle.schedule(shard);
-        }
-    }
-
     fn shard_ids(&self) -> Vec<ContractAddress> {
         self.manager.shard_ids()
     }

crates/node/sharding/src/lib.rs

@@ -35,9 +35,6 @@ pub trait ShardProvider: Send + Sync + 'static {
     /// Resolve a shard's API by ID (for read operations).
     fn starknet_api(&self, shard_id: ContractAddress) -> Result<Self::Api, ErrorObjectOwned>;
 
-    /// Schedule a shard for execution after a write operation.
-    fn schedule(&self, shard_id: ContractAddress);
-
     /// List all registered shard IDs.
     fn shard_ids(&self) -> Vec<ContractAddress>;
 
@@ -211,10 +208,7 @@ impl<P: ShardProvider> ShardApiServer for ShardRpc<P> {
         invoke_transaction: BroadcastedInvokeTx,
     ) -> RpcResult<AddInvokeTransactionResponse> {
         let api = self.provider.starknet_api(shard_id)?;
-        let result =
-            StarknetWriteApiServer::add_invoke_transaction(&api, invoke_transaction).await?;
-        self.provider.schedule(shard_id);
-        Ok(result)
+        StarknetWriteApiServer::add_invoke_transaction(&api, invoke_transaction).await
     }
 
     async fn add_declare_transaction(
@@ -223,10 +217,7 @@ impl<P: ShardProvider> ShardApiServer for ShardRpc<P> {
         declare_transaction: BroadcastedDeclareTx,
     ) -> RpcResult<AddDeclareTransactionResponse> {
         let api = self.provider.starknet_api(shard_id)?;
-        let result =
-            StarknetWriteApiServer::add_declare_transaction(&api, declare_transaction).await?;
-        self.provider.schedule(shard_id);
-        Ok(result)
+        StarknetWriteApiServer::add_declare_transaction(&api, declare_transaction).await
     }
 
     async fn add_deploy_account_transaction(
@@ -235,13 +226,8 @@ impl<P: ShardProvider> ShardApiServer for ShardRpc<P> {
         deploy_account_transaction: BroadcastedDeployAccountTx,
     ) -> RpcResult<AddDeployAccountTransactionResponse> {
         let api = self.provider.starknet_api(shard_id)?;
-        let result = StarknetWriteApiServer::add_deploy_account_transaction(
-            &api,
-            deploy_account_transaction,
-        )
-        .await?;
-        self.provider.schedule(shard_id);
-        Ok(result)
+        StarknetWriteApiServer::add_deploy_account_transaction(&api, deploy_account_transaction)
+            .await
     }
 
     async fn trace_transaction(

crates/rpc/rpc-server/src/shard.rs

@@ -12,6 +12,7 @@ katana-db.workspace = true
 katana-executor.workspace = true
 katana-gas-price-oracle.workspace = true
 katana-pool.workspace = true
+katana-pool-api.workspace = true
 katana-primitives.workspace = true
 katana-provider.workspace = true
 katana-rpc-client.workspace = true
@@ -21,6 +22,7 @@ katana-rpc-api.workspace = true
 katana-tasks.workspace = true
 
 anyhow.workspace = true
+futures.workspace = true
 parking_lot.workspace = true
 tokio.workspace = true
 tracing.workspace = true

crates/sharding/Cargo.toml

@@ -1,4 +1,5 @@
 pub mod manager;
+pub mod pool;
 pub mod runtime;
 pub mod scheduler;
 pub mod shard;

crates/sharding/src/lib.rs

@@ -1,4 +1,3 @@
-use std::collections::HashMap;
 use std::fmt::Debug;
 use std::sync::Arc;
 
@@ -12,9 +11,9 @@ use katana_rpc_client::starknet::Client as StarknetClient;
 use katana_rpc_server::starknet::StarknetApiConfig;
 use katana_rpc_types::block::GetBlockWithTxHashesResponse;
 use katana_tasks::TaskSpawner;
-use parking_lot::RwLock;
 use url::Url;
 
+use crate::scheduler::{Scheduler, ShardRegistry};
 use crate::shard::{Shard, ShardId};
 
 type InitialBlockEnvFetcher = dyn Fn() -> Result<BlockEnv> + Send + Sync + 'static;
@@ -119,7 +118,8 @@ pub struct LazyShardManager {
 }
 
 struct LazyShardManagerInner {
-    shards: RwLock<HashMap<ShardId, Arc<Shard>>>,
+    shards: ShardRegistry,
+    scheduler: Scheduler,
     // Shared resources for lazy shard creation
     chain_spec: Arc<ChainSpec>,
     executor_factory: Arc<dyn ExecutorFactory>,
@@ -136,6 +136,7 @@ impl LazyShardManager {
         gas_oracle: GasPriceOracle,
         starknet_api_config: StarknetApiConfig,
         task_spawner: TaskSpawner,
+        scheduler: Scheduler,
         base_chain_url: Url,
     ) -> Self {
         let initial_block_env_fetcher = new_base_chain_block_env_fetcher(base_chain_url);
@@ -146,6 +147,7 @@ impl LazyShardManager {
             gas_oracle,
             starknet_api_config,
             task_spawner,
+            scheduler,
             initial_block_env_fetcher,
         )
     }
@@ -156,11 +158,13 @@ impl LazyShardManager {
         gas_oracle: GasPriceOracle,
         starknet_api_config: StarknetApiConfig,
         task_spawner: TaskSpawner,
+        scheduler: Scheduler,
         initial_block_env_fetcher: Arc<InitialBlockEnvFetcher>,
     ) -> Self {
         Self {
             inner: Arc::new(LazyShardManagerInner {
-                shards: RwLock::new(HashMap::new()),
+                shards: scheduler.shard_registry(),
+                scheduler,
                 chain_spec,
                 executor_factory,
                 gas_oracle,
@@ -212,9 +216,11 @@ impl ShardManager for LazyShardManager {
             self.inner.starknet_api_config.clone(),
             self.inner.task_spawner.clone(),
             initial_block_env,
+            self.inner.scheduler.clone(),
         )?);
 
         shards.insert(id, Arc::clone(&shard));
+
         Ok(shard)
     }
 

crates/sharding/src/manager.rs

@@ -0,0 +1,93 @@
+use std::fmt;
+use std::future::Future;
+use std::sync::Arc;
+
+use futures::channel::mpsc::Receiver;
+use katana_pool::{PendingTransactions, TransactionPool, TxPool};
+use katana_pool_api::PoolResult;
+use katana_primitives::contract::Nonce;
+use katana_primitives::transaction::TxHash;
+use katana_primitives::ContractAddress;
+
+use crate::scheduler::Scheduler;
+use crate::shard::ShardId;
+
+/// Transaction pool wrapper that triggers shard scheduling when a tx is accepted.
+#[derive(Clone)]
+pub struct ShardPool {
+    inner: TxPool,
+    scheduler: Scheduler,
+    shard_id: ShardId,
+}
+
+impl ShardPool {
+    pub fn new(inner: TxPool, scheduler: Scheduler, shard_id: ShardId) -> Self {
+        Self { inner, scheduler, shard_id }
+    }
+}
+
+impl TransactionPool for ShardPool {
+    type Transaction = <TxPool as TransactionPool>::Transaction;
+    type Ordering = <TxPool as TransactionPool>::Ordering;
+    type Validator = <TxPool as TransactionPool>::Validator;
+
+    fn add_transaction(
+        &self,
+        tx: Self::Transaction,
+    ) -> impl Future<Output = PoolResult<TxHash>> + Send {
+        let inner = self.inner.clone();
+        let scheduler = self.scheduler.clone();
+        let shard_id = self.shard_id;
+
+        async move {
+            let result = inner.add_transaction(tx).await;
+
+            if result.is_ok() {
+                scheduler.schedule(shard_id);
+            }
+
+            result
+        }
+    }
+
+    fn pending_transactions(&self) -> PendingTransactions<Self::Transaction, Self::Ordering> {
+        self.inner.pending_transactions()
+    }
+
+    fn contains(&self, hash: TxHash) -> bool {
+        self.inner.contains(hash)
+    }
+
+    fn get(&self, hash: TxHash) -> Option<Arc<Self::Transaction>> {
+        self.inner.get(hash)
+    }
+
+    fn add_listener(&self) -> Receiver<TxHash> {
+        self.inner.add_listener()
+    }
+
+    fn remove_transactions(&self, hashes: &[TxHash]) {
+        self.inner.remove_transactions(hashes);
+    }
+
+    fn size(&self) -> usize {
+        self.inner.size()
+    }
+
+    fn validator(&self) -> &Self::Validator {
+        self.inner.validator()
+    }
+
+    fn get_nonce(&self, address: ContractAddress) -> Option<Nonce> {
+        self.inner.get_nonce(address)
+    }
+}
+
+impl fmt::Debug for ShardPool {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ShardPool")
+            .field("size", &self.inner.size())
+            .field("shard_id", &self.shard_id)
+            .finish()
+    }
+}

crates/sharding/src/pool.rs

@@ -1,11 +1,10 @@
-use std::sync::Arc;
 use std::thread;
 use std::time::{Duration, Instant};
 
 use anyhow::Result;
 
 use crate::scheduler::Scheduler;
-use crate::shard::Shard;
+use crate::shard::ShardId;
 use crate::worker::Worker;
 
 /// Owns the execution resources (scheduler + worker threads) for the shard node.
@@ -41,8 +40,8 @@ pub struct RuntimeHandle {
 
 impl RuntimeHandle {
     /// Schedule a shard for execution.
-    pub fn schedule(&self, shard: Arc<Shard>) {
-        self.scheduler.schedule(shard);
+    pub fn schedule(&self, id: ShardId) {
+        self.scheduler.schedule(id);
     }
 
     /// Returns a reference to the underlying scheduler.