apollo_gateway: use get_nonce outside extract_state_nonce_and_run_validations

Author: itamar-starkware

crates/apollo_gateway/src/gateway.rs

@@ -41,7 +41,7 @@ use crate::errors::{
     GatewayResult,
 };
 use crate::metrics::{register_metrics, GatewayMetricHandle, GATEWAY_ADD_TX_LATENCY};
-use crate::state_reader::StateReaderFactory;
+use crate::state_reader::{MempoolStateReader, StateReaderFactory};
 use crate::stateful_transaction_validator::{
     StatefulTransactionValidatorFactory,
     StatefulTransactionValidatorFactoryTrait,
@@ -151,7 +151,7 @@ impl Gateway {
                 transaction_converter_err_to_deprecated_gw_err(&tx_signature, e)
             })?;
 
-        let blocking_task =
+        let mut blocking_task =
             ProcessTxBlockingTask::new(self, executable_tx, tokio::runtime::Handle::current())
                 .await
                 .map_err(|e| {
@@ -162,10 +162,22 @@ impl Gateway {
                     metric_counters.record_add_tx_failure(&e);
                     e
                 })?;
+
+        let state_reader = blocking_task.get_state_reader().await?;
+
+        let account_address = blocking_task.executable_tx.contract_address();
+        let account_nonce =
+            state_reader.get_account_nonce(account_address).await.map_err(|e| {
+                StarknetError::internal_with_logging("Failed to get account nonce", e)
+            })?;
+
+        let stateful_tx_validator = blocking_task.create_stateful_validator(state_reader).await?;
+
         // Run the blocking task in the current span.
         let curr_span = Span::current();
-        let handle =
-            tokio::task::spawn_blocking(move || curr_span.in_scope(|| blocking_task.process_tx()));
+        let handle = tokio::task::spawn_blocking(move || {
+            curr_span.in_scope(|| blocking_task.process_tx(account_nonce, stateful_tx_validator))
+        });
         let handle_result = handle.await;
         let nonce = match handle_result {
             Ok(Ok(nonce)) => nonce,
@@ -240,7 +252,8 @@ impl Gateway {
 /// CPU-intensive transaction processing, spawned in a blocking thread to avoid blocking other tasks
 /// from running.
 struct ProcessTxBlockingTask {
-    stateful_tx_validator: Box<dyn StatefulTransactionValidatorTrait + Send>,
+    state_reader_factory: Arc<dyn StateReaderFactory>,
+    stateful_tx_validator_factory: Arc<dyn StatefulTransactionValidatorFactoryTrait>,
     mempool_client: SharedMempoolClient,
     executable_tx: AccountTransaction,
     runtime: tokio::runtime::Handle,
@@ -252,21 +265,52 @@ impl ProcessTxBlockingTask {
         executable_tx: AccountTransaction,
         runtime: tokio::runtime::Handle,
     ) -> GatewayResult<Self> {
-        let stateful_tx_validator = gateway
-            .stateful_tx_validator_factory
-            .instantiate_validator(gateway.state_reader_factory.as_ref())
-            .await?;
         Ok(Self {
-            stateful_tx_validator,
+            state_reader_factory: gateway.state_reader_factory.clone(),
+            stateful_tx_validator_factory: gateway.stateful_tx_validator_factory.clone(),
             mempool_client: gateway.mempool_client.clone(),
             executable_tx,
             runtime,
         })
     }
 
-    fn process_tx(mut self) -> GatewayResult<Nonce> {
-        let nonce = self.stateful_tx_validator.extract_state_nonce_and_run_validations(
+    pub async fn get_state_reader(&self) -> GatewayResult<Box<dyn MempoolStateReader>> {
+        let state_reader = self
+        .stateful_tx_validator_factory
+        .get_state_reader_for_validation(self.state_reader_factory.as_ref())
+        .await
+        .map_err(|e| {
+            StarknetError::internal_with_logging(
+                "Failed to get state reader from latest block",
+                e,
+            )
+        })?;
+        Ok(state_reader)
+    }
+
+    pub async fn create_stateful_validator(
+        &mut self,
+        state_reader: Box<dyn MempoolStateReader>,
+    ) -> GatewayResult<Box<dyn StatefulTransactionValidatorTrait>> {
+        self.stateful_tx_validator_factory
+            .create_validator_from_state_reader(state_reader)
+            .await
+            .map_err(|e| {
+                StarknetError::internal_with_logging(
+                    "Failed to create stateful validator from state reader",
+                    e,
+                )
+            })
+    }
+
+    fn process_tx(
+        self,
+        account_nonce: Nonce,
+        mut stateful_tx_validator: Box<dyn StatefulTransactionValidatorTrait>,
+    ) -> GatewayResult<Nonce> {
+        let nonce = stateful_tx_validator.extract_state_nonce_and_run_validations(
             &self.executable_tx,
+            account_nonce,
             self.mempool_client,
             self.runtime,
         )?;

crates/apollo_gateway/src/gateway_test.rs

@@ -97,8 +97,6 @@ use crate::state_reader_test_utils::{local_test_state_reader_factory, TestStateR
 use crate::stateful_transaction_validator::{
     MockStatefulTransactionValidatorFactoryTrait,
     MockStatefulTransactionValidatorTrait,
-    StatefulTransactionValidatorFactoryTrait,
-    StatefulTransactionValidatorTrait,
 };
 use crate::stateless_transaction_validator::MockStatelessTransactionValidatorTrait;
 
@@ -316,17 +314,12 @@ async fn run_add_tx_and_extract_metrics(
 }
 
 async fn process_tx_task(
-    stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
+    stateful_tx_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
 ) -> ProcessTxBlockingTask {
     let state_reader_factory = Arc::new(MockStateReaderFactory::new());
-    let stateful_tx_validator = StatefulTransactionValidatorFactoryTrait::instantiate_validator(
-        &stateful_transaction_validator_factory,
-        state_reader_factory.as_ref(),
-    )
-    .await
-    .expect("instantiate_validator should be mocked in tests");
     ProcessTxBlockingTask {
-        stateful_tx_validator,
+        state_reader_factory,
+        stateful_tx_validator_factory: Arc::new(stateful_tx_validator_factory),
         mempool_client: Arc::new(MockMempoolClient::new()),
         executable_tx: executable_invoke_tx(invoke_args()),
         runtime: tokio::runtime::Handle::current(),
@@ -558,7 +551,7 @@ fn test_full_cycle_dump_deserialize_authorized_declarer_accounts(
 async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_fails(
     #[case] error_code: StarknetErrorCode,
     mut mock_stateful_transaction_validator: MockStatefulTransactionValidatorTrait,
-    mut mock_stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
+    mock_stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
 ) {
     let expected_error = StarknetError {
         code: error_code.clone(),
@@ -567,19 +560,15 @@ async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_f
 
     mock_stateful_transaction_validator
         .expect_extract_state_nonce_and_run_validations()
-        .return_once(|_, _, _| Err(expected_error));
-
-    mock_stateful_transaction_validator_factory.expect_instantiate_validator().return_once(|_| {
-        Box::pin(async {
-            Ok::<Box<dyn StatefulTransactionValidatorTrait>, _>(Box::new(
-                mock_stateful_transaction_validator,
-            ))
-        })
-    });
+        .return_once(|_, _, _, _| Err(expected_error));
 
     let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory).await;
-
-    let result = tokio::task::spawn_blocking(move || process_tx_task.process_tx()).await.unwrap();
+    let account_nonce = nonce!(0);
+    let result = tokio::task::spawn_blocking(move || {
+        process_tx_task.process_tx(account_nonce, Box::new(mock_stateful_transaction_validator))
+    })
+    .await
+    .unwrap();
 
     assert!(result.is_err());
     assert_eq!(result.unwrap_err().code, error_code);
@@ -620,7 +609,7 @@ async fn add_tx_returns_error_when_instantiating_validator_fails(
     let error_code = StarknetErrorCode::UnknownErrorCode("StarknetErrorCode.InternalError".into());
     let expected_error = StarknetError { code: error_code.clone(), message: "placeholder".into() };
     mock_stateful_transaction_validator_factory
-        .expect_instantiate_validator()
+        .expect_get_state_reader_for_validation()
         .return_once(|_| Box::pin(async { Err::<_, _>(expected_error) }));
 
     // Build gateway and inject the failing factory.

crates/apollo_gateway/src/rpc_state_reader.rs

@@ -111,6 +111,22 @@ impl MempoolStateReader for RpcStateReader {
         let block_info = block_header.try_into()?;
         Ok(block_info)
     }
+
+    async fn get_account_nonce(
+        &self,
+        contract_address: ContractAddress,
+    ) -> StateResult<Nonce> {
+        let get_nonce_params = GetNonceParams { block_id: self.block_id, contract_address };
+        let result = self.send_rpc_request("starknet_getNonce", get_nonce_params);
+        match result {
+            Ok(value) => {
+                let nonce: Nonce = serde_json::from_value(value).map_err(serde_err_to_state_err)?;
+                Ok(nonce)
+            }
+            Err(RPCStateReaderError::ContractAddressNotFound(_)) => Ok(Nonce::default()),
+            Err(e) => Err(e)?,
+        }
+    }
 }
 
 impl BlockifierStateReader for RpcStateReader {

crates/apollo_gateway/src/state_reader.rs

@@ -13,6 +13,10 @@ use starknet_types_core::felt::Felt;
 #[async_trait]
 pub trait MempoolStateReader: BlockifierStateReader + Send + Sync {
     async fn get_block_info(&self) -> Result<BlockInfo, StateError>;
+    async fn get_account_nonce(
+        &self,
+        contract_address: ContractAddress,
+    ) -> Result<Nonce, StateError>;
 }
 
 #[async_trait]
@@ -32,6 +36,13 @@ impl MempoolStateReader for Box<dyn MempoolStateReader> {
     async fn get_block_info(&self) -> Result<BlockInfo, StateError> {
         self.as_ref().get_block_info().await
     }
+    async fn get_account_nonce(
+        &self,
+        contract_address: ContractAddress,
+    ) -> Result<Nonce, StateError> {
+        crate::state_reader::MempoolStateReader::get_account_nonce(self.as_ref(), contract_address)
+            .await
+    }
 }
 
 impl BlockifierStateReader for Box<dyn MempoolStateReader> {

crates/apollo_gateway/src/state_reader_test_utils.rs

@@ -28,6 +28,13 @@ impl MempoolStateReader for TestStateReader {
     async fn get_block_info(&self) -> Result<BlockInfo, StateError> {
         Ok(self.block_info.clone())
     }
+
+    async fn get_account_nonce(
+        &self,
+        contract_address: ContractAddress,
+    ) -> Result<Nonce, StateError> {
+        self.blockifier_state_reader.get_nonce_at(contract_address)
+    }
 }
 
 impl BlockifierStateReader for TestStateReader {

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -47,11 +47,12 @@ pub trait StatefulTransactionValidatorFactoryTrait: Send + Sync {
         state_reader_factory: &dyn StateReaderFactory,
     ) -> StatefulTransactionValidatorResult<Box<dyn StatefulTransactionValidatorTrait>>;
 
+    // 1) Get a state reader for validation (async, uses the factory).
     async fn get_state_reader_for_validation(
         &self,
         state_reader_factory: &dyn StateReaderFactory,
     ) -> StatefulTransactionValidatorResult<Box<dyn MempoolStateReader>>;
-
+    // 2) Create the validator from a provided state reader (async, uses ChainInfo/config).
     async fn create_validator_from_state_reader(
         &self,
         state_reader: Box<dyn MempoolStateReader>,
@@ -132,6 +133,7 @@ pub trait StatefulTransactionValidatorTrait: Send {
     fn extract_state_nonce_and_run_validations(
         &mut self,
         executable_tx: &ExecutableTransaction,
+        account_nonce: Nonce,
         mempool_client: SharedMempoolClient,
         runtime: tokio::runtime::Handle,
     ) -> StatefulTransactionValidatorResult<Nonce>;
@@ -153,11 +155,10 @@ impl<B: BlockifierStatefulValidatorTrait + Send> StatefulTransactionValidatorTra
     fn extract_state_nonce_and_run_validations(
         &mut self,
         executable_tx: &ExecutableTransaction,
+        account_nonce: Nonce,
         mempool_client: SharedMempoolClient,
         runtime: tokio::runtime::Handle,
     ) -> StatefulTransactionValidatorResult<Nonce> {
-        let address = executable_tx.contract_address();
-        let account_nonce = self.get_nonce(address)?;
         self.run_transaction_validations(executable_tx, account_nonce, mempool_client, runtime)?;
         Ok(account_nonce)
     }

crates/apollo_gateway/src/stateful_transaction_validator_test.rs

@@ -64,32 +64,24 @@ async fn test_get_nonce_fail_on_extract_state_nonce_and_run_validations() {
         // The mempool does not have any transactions from the sender.
         Ok(false)
     });
-    let mempool_client = Arc::new(mock_mempool_client);
-    let runtime = tokio::runtime::Handle::current();
+    // Removed unused mempool client and runtime in this test since nonce is fetched outside now.
 
     let mut stateful_validator = StatefulTransactionValidator {
         config: StatefulTransactionValidatorConfig::default(),
         blockifier_stateful_tx_validator: mock_blockifier_validator,
     };
 
     let executable_tx = create_executable_invoke_tx(CairoVersion::Cairo1(RunnableCairo1::Casm));
-    let result = tokio::task::spawn_blocking(move || {
-        stateful_validator.extract_state_nonce_and_run_validations(
-            &executable_tx,
-            mempool_client,
-            runtime,
-        )
+    // get_nonce should fail and return an InternalError.
+    let err = tokio::task::spawn_blocking(move || {
+        stateful_validator.get_nonce(executable_tx.contract_address())
     })
     .await
-    .unwrap();
+    .unwrap()
+    .unwrap_err();
     assert_eq!(
-        result,
-        Err(StarknetError {
-            code: StarknetErrorCode::UnknownErrorCode(
-                "StarknetErrorCode.InternalError".to_string()
-            ),
-            message: "Internal error".to_string(),
-        })
+        err.code,
+        StarknetErrorCode::UnknownErrorCode("StarknetErrorCode.InternalError".to_string())
     );
 }
 
@@ -147,6 +139,7 @@ async fn test_extract_state_nonce_and_run_validations(
     let result = tokio::task::spawn_blocking(move || {
         stateful_validator.extract_state_nonce_and_run_validations(
             &executable_tx,
+            account_nonce,
             mempool_client,
             runtime,
         )

crates/apollo_gateway/src/sync_state_reader.rs

@@ -88,6 +88,18 @@ impl MempoolStateReader for SyncStateReader {
 
         Ok(block_info)
     }
+
+    async fn get_account_nonce(&self, contract_address: ContractAddress) -> StateResult<Nonce> {
+        let res = self.state_sync_client.get_nonce_at(self.block_number, contract_address).await;
+
+        match res {
+            Ok(value) => Ok(value),
+            Err(StateSyncClientError::StateSyncError(StateSyncError::ContractNotFound(_))) => {
+                Ok(Nonce::default())
+            }
+            Err(e) => Err(StateError::StateReadError(e.to_string())),
+        }
+    }
 }
 
 impl BlockifierStateReader for SyncStateReader {