refactor(rpc): pool abstraction (#348)

Author: kariy

Cargo.lock

@@ -10,6 +10,7 @@ use katana_primitives::transaction::ExecutableTxWithHash;
 use katana_primitives::utils::transaction::compute_invoke_v3_tx_hash;
 use katana_primitives::Felt;
 use katana_rpc_types::broadcasted::BroadcastedInvokeTx;
+use katana_rpc_types::{BroadcastedTx, BroadcastedTxWithChainId};
 use starknet::accounts::{Account, ExecutionEncoder, ExecutionEncoding, SingleOwnerAccount};
 use starknet::core::types::{BlockId, BlockTag, Call};
 use starknet::macros::{felt, selector};
@@ -119,7 +120,7 @@ pub fn invoke_executable_tx(
         broadcasted_tx.signature = vec![]
     }
 
-    ExecutableTxWithHash::new(broadcasted_tx.into_inner(chain_id).into())
+    BroadcastedTxWithChainId { tx: BroadcastedTx::Invoke(broadcasted_tx), chain: chain_id }.into()
 }
 
 #[rstest::fixture]

crates/executor/tests/fixtures/transaction.rs

@@ -11,6 +11,7 @@ katana-executor.workspace = true
 katana-core.workspace = true
 katana-primitives.workspace = true
 katana-metrics.workspace = true
+katana-pool.workspace = true
 katana-rpc.workspace = true
 katana-rpc-api.workspace = true
 katana-provider-api.workspace = true

crates/gateway/gateway-server/Cargo.toml

@@ -7,6 +7,7 @@ use katana_gateway_types::{
     Block, ConfirmedReceipt, ConfirmedTransaction, ContractClass, ErrorCode, GatewayError,
     ReceiptBody, StateUpdate, StateUpdateWithBlock,
 };
+use katana_pool::TxPool;
 use katana_primitives::block::{BlockHash, BlockIdOrTag, BlockNumber};
 use katana_primitives::class::{ClassHash, CompiledClass, ContractClassCompilationError};
 use katana_provider_api::block::{BlockIdReader, BlockProvider, BlockStatusProvider};
@@ -20,7 +21,7 @@ use starknet::core::types::ResourcePrice;
 /// Shared application state containing the backend
 #[derive(Clone)]
 pub struct AppState {
-    pub api: StarknetApi<BlockifierFactory, BlockProducer<BlockifierFactory>>,
+    pub api: StarknetApi<BlockifierFactory, TxPool, BlockProducer<BlockifierFactory>>,
 }
 
 impl AppState {

crates/gateway/gateway-server/src/handlers.rs

@@ -75,13 +75,18 @@ pub struct GatewayServer {
     health_check: bool,
     metered: bool,
 
-    starknet_api: StarknetApi<BlockifierFactory, BlockProducer<BlockifierFactory>>,
+    starknet_api:
+        StarknetApi<BlockifierFactory, katana_pool::TxPool, BlockProducer<BlockifierFactory>>,
 }
 
 impl GatewayServer {
     /// Create a new feeder gateway server.
     pub fn new(
-        starknet_api: StarknetApi<BlockifierFactory, BlockProducer<BlockifierFactory>>,
+        starknet_api: StarknetApi<
+            BlockifierFactory,
+            katana_pool::TxPool,
+            BlockProducer<BlockifierFactory>,
+        >,
     ) -> Self {
         Self {
             timeout: DEFAULT_GATEWAY_TIMEOUT,

crates/gateway/gateway-server/src/lib.rs

@@ -30,7 +30,7 @@ pub enum PoolError {
 pub type PoolResult<T> = Result<T, PoolError>;
 
 /// Represents a complete transaction pool.
-pub trait TransactionPool {
+pub trait TransactionPool: Send + Sync {
     /// The pool's transaction type.
     type Transaction: PoolTransaction;
 
@@ -67,6 +67,13 @@ pub trait TransactionPool {
 
     /// Get a reference to the pool's validator.
     fn validator(&self) -> &Self::Validator;
+
+    /// Get the next expected nonce for an account based on pending transactions in the pool.
+    ///
+    /// Returns `Some(nonce)` if there are pending transactions for this account,
+    /// where nonce is the next expected nonce (highest pending nonce + 1).
+    /// Returns `None` if no pending transactions exist for this account.
+    fn get_nonce(&self, address: ContractAddress) -> Option<Nonce>;
 }
 
 // the transaction type is recommended to implement a cheap clone (eg ref-counting) so that it

crates/pool/pool-api/src/lib.rs

@@ -9,7 +9,9 @@ use katana_pool_api::{
     PendingTransactions, PendingTx, PoolError, PoolOrd, PoolResult, PoolTransaction, Subscription,
     TransactionPool, TxId,
 };
+use katana_primitives::contract::Nonce;
 use katana_primitives::transaction::TxHash;
+use katana_primitives::ContractAddress;
 use parking_lot::RwLock;
 use tokio::sync::mpsc;
 use tracing::{error, trace, warn, Instrument};
@@ -227,6 +229,17 @@ where
     fn validator(&self) -> &Self::Validator {
         &self.inner.validator
     }
+
+    fn get_nonce(&self, address: ContractAddress) -> Option<Nonce> {
+        self.inner
+            .transactions
+            .read()
+            .iter()
+            .filter(|tx| tx.tx.sender() == address)
+            .map(|tx| tx.tx.nonce())
+            .max()
+            .map(|max_nonce| max_nonce + 1)
+    }
 }
 
 impl<T, V, O> Clone for Pool<T, V, O>
@@ -494,4 +507,97 @@ mod tests {
     #[test]
     #[ignore = "Txs dependency management not fully implemented yet"]
     fn dependent_txs_random_insertion() {}
+
+    #[tokio::test]
+    async fn get_nonce_returns_none_for_unknown_address() {
+        let pool = TestPool::test();
+        let unknown_address = ContractAddress::from(Felt::from_hex("0xdead").unwrap());
+        assert_eq!(pool.get_nonce(unknown_address), None);
+    }
+
+    #[tokio::test]
+    async fn get_nonce_returns_next_nonce_for_single_pending_tx() {
+        let pool = TestPool::test();
+        let sender = ContractAddress::from(Felt::from_hex("0x1337").unwrap());
+        let nonce = Nonce::from(5u128);
+
+        let tx = PoolTx::new().with_sender(sender).with_nonce(nonce);
+        pool.add_transaction(tx).await.unwrap();
+
+        // Should return nonce + 1
+        assert_eq!(pool.get_nonce(sender), Some(Nonce::from(6u128)));
+    }
+
+    #[tokio::test]
+    async fn get_nonce_returns_highest_nonce_plus_one_for_multiple_txs() {
+        let pool = TestPool::test();
+        let sender = ContractAddress::from(Felt::from_hex("0x1337").unwrap());
+
+        // Add transactions with nonces 1, 5, 3 (out of order)
+        let tx1 = PoolTx::new().with_sender(sender).with_nonce(Nonce::from(1u128));
+        let tx2 = PoolTx::new().with_sender(sender).with_nonce(Nonce::from(3u128));
+        let tx3 = PoolTx::new().with_sender(sender).with_nonce(Nonce::from(2u128));
+
+        pool.add_transaction(tx1).await.unwrap();
+        pool.add_transaction(tx2).await.unwrap();
+        pool.add_transaction(tx3).await.unwrap();
+
+        // Should return max(1, 3, 2) + 1 = 4
+        assert_eq!(pool.get_nonce(sender), Some(Nonce::from(4u128)));
+    }
+
+    #[tokio::test]
+    async fn get_nonce_isolated_per_address() {
+        let pool = TestPool::test();
+
+        let sender1 = ContractAddress::from(Felt::from_hex("0x1").unwrap());
+        let sender2 = ContractAddress::from(Felt::from_hex("0x2").unwrap());
+        let sender3 = ContractAddress::from(Felt::from_hex("0x3").unwrap());
+
+        // Add transactions from different senders
+        let tx1 = PoolTx::new().with_sender(sender1).with_nonce(Nonce::from(10u128));
+        let tx2 = PoolTx::new().with_sender(sender2).with_nonce(Nonce::from(20u128));
+        let tx3 = PoolTx::new().with_sender(sender1).with_nonce(Nonce::from(15u128));
+
+        pool.add_transaction(tx1).await.unwrap();
+        pool.add_transaction(tx2).await.unwrap();
+        pool.add_transaction(tx3).await.unwrap();
+
+        // Each sender should have their own max nonce
+        assert_eq!(pool.get_nonce(sender1), Some(Nonce::from(16u128))); // max(10, 15) + 1
+        assert_eq!(pool.get_nonce(sender2), Some(Nonce::from(21u128))); // 20 + 1
+        assert_eq!(pool.get_nonce(sender3), None); // No txs from sender3
+    }
+
+    #[tokio::test]
+    async fn get_nonce_updates_after_transaction_removal() {
+        let pool = TestPool::test();
+        let sender = ContractAddress::from(Felt::from_hex("0x1337").unwrap());
+
+        let tx1 = PoolTx::new().with_sender(sender).with_nonce(Nonce::from(1u128));
+        let tx2 = PoolTx::new().with_sender(sender).with_nonce(Nonce::from(2u128));
+        let tx3 = PoolTx::new().with_sender(sender).with_nonce(Nonce::from(3u128));
+
+        let hash1 = tx1.hash();
+        let hash3 = tx3.hash();
+
+        pool.add_transaction(tx1).await.unwrap();
+        pool.add_transaction(tx2.clone()).await.unwrap();
+        pool.add_transaction(tx3).await.unwrap();
+
+        // Should be 4 (max nonce 3 + 1)
+        assert_eq!(pool.get_nonce(sender), Some(Nonce::from(4u128)));
+
+        // Remove transactions with nonce 1 and 3
+        pool.remove_transactions(&[hash1, hash3]);
+
+        // Should now be 3 (only tx with nonce 2 remains)
+        assert_eq!(pool.get_nonce(sender), Some(Nonce::from(3u128)));
+
+        // Remove last transaction
+        pool.remove_transactions(&[tx2.hash()]);
+
+        // Should be None (no transactions left)
+        assert_eq!(pool.get_nonce(sender), None);
+    }
 }

crates/pool/pool/src/pool.rs

@@ -27,7 +27,6 @@ use katana_primitives::env::{BlockEnv, CfgEnv};
 use katana_primitives::transaction::{ExecutableTx, ExecutableTxWithHash};
 use katana_primitives::Felt;
 use katana_provider::api::state::StateProvider;
-use katana_provider::api::ProviderError;
 use parking_lot::Mutex;
 
 use super::ValidationResult;
@@ -72,19 +71,6 @@ impl TxValidator {
         this.block_env = block_env;
         this.state = Arc::new(new_state);
     }
-
-    // NOTE:
-    // If you check the get_nonce method of StatefulValidator in blockifier, under the hood it
-    // unwraps the Option to get the state of the TransactionExecutor struct. StatefulValidator
-    // guaranteees that the state will always be present so it is safe to uwnrap. However, this
-    // safety is not guaranteed by TransactionExecutor itself.
-    pub fn pool_nonce(&self, address: ContractAddress) -> Result<Option<Nonce>, ProviderError> {
-        let this = self.inner.lock();
-        match this.pool_nonces.get(&address) {
-            Some(nonce) => Ok(Some(*nonce)),
-            None => Ok(this.state.nonce(address)?),
-        }
-    }
 }
 
 impl Debug for Inner {

crates/pool/pool/src/validation/stateful.rs

@@ -62,7 +62,7 @@ impl ContractClass {
                     .map(|f| f.value.clone().into())
                     .collect::<Vec<Felt>>();
 
-                compute_sierra_class_hash(&abi_str, &class.entry_points_by_type, sierra_program)
+                Ok(compute_sierra_class_hash(&abi_str, &class.entry_points_by_type, sierra_program))
             }
 
             Self::Legacy(class) => compute_legacy_class_hash(class),
@@ -232,7 +232,7 @@ pub fn compute_sierra_class_hash(
     abi: &str,
     entry_points_by_type: &ContractEntryPoints,
     sierra_program: &[Felt],
-) -> Result<Felt, ComputeClassHashError> {
+) -> Felt {
     let mut hasher = starknet_crypto::PoseidonHasher::new();
     hasher.update(short_string!("CONTRACT_CLASS_V0.1.0"));
 
@@ -245,7 +245,7 @@ pub fn compute_sierra_class_hash(
     // Hashes Sierra program
     hasher.update(Poseidon::hash_array(sierra_program));
 
-    Ok(normalize_address(hasher.finalize()))
+    normalize_address(hasher.finalize())
 }
 
 /// Computes the hash of a legacy contract class.