apollo_gateway: use get_nonce outside extract_state_nonce_and_run_validations

Author: itamar-starkware

crates/apollo_gateway/src/gateway.rs

@@ -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,23 @@ impl Gateway {
                     metric_counters.record_add_tx_failure(&e);
                     e
                 })?;
+
+        // Obtain account nonce before spawning the blocking computation.
+        let account_nonce = blocking_task.get_nonce().map_err(|e| {
+            info!(
+                "Gateway validation failed for tx with signature: {:?} with error: {}",
+                tx.signature(),
+                e
+            );
+            metric_counters.record_add_tx_failure(&e);
+            e
+        })?;
+
         // 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))
+        });
         let handle_result = handle.await;
         let nonce = match handle_result {
             Ok(Ok(nonce)) => nonce,
@@ -264,9 +277,15 @@ impl ProcessTxBlockingTask {
         })
     }
 
-    fn process_tx(mut self) -> GatewayResult<Nonce> {
+    fn get_nonce(&mut self) -> GatewayResult<Nonce> {
+        let address = self.executable_tx.contract_address();
+        self.stateful_tx_validator.get_nonce(address)
+    }
+
+    fn process_tx(mut self, account_nonce: Nonce) -> GatewayResult<Nonce> {
         let nonce = self.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

@@ -567,7 +567,7 @@ 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));
+        .return_once(|_, _, _, _| Err(expected_error));
 
     mock_stateful_transaction_validator_factory.expect_instantiate_validator().return_once(|_| {
         Box::pin(async {
@@ -578,8 +578,10 @@ async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_f
     });
 
     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))
+        .await
+        .unwrap();
 
     assert!(result.is_err());
     assert_eq!(result.unwrap_err().code, error_code);

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -107,6 +107,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>;
@@ -128,11 +129,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,
         )