apollo_gateway: extract async mempool query from blocking task

itamar-starkware · itamar-starkware · commit f8cf51905a14 · 2025-11-06T17:28:06.000+02:00
diff --git a/crates/apollo_gateway/src/gateway.rs b/crates/apollo_gateway/src/gateway.rs
@@ -19,7 +19,11 @@ use apollo_gateway_types::gateway_types::{
     InvokeGatewayOutput,
 };
 use apollo_infra::component_definitions::ComponentStarter;
-use apollo_mempool_types::communication::{AddTransactionArgsWrapper, SharedMempoolClient};
+use apollo_mempool_types::communication::{
+    AddTransactionArgsWrapper,
+    MempoolClientError,
+    SharedMempoolClient,
+};
 use apollo_mempool_types::mempool_types::AddTransactionArgs;
 use apollo_network_types::network_types::BroadcastedMessageMetadata;
 use apollo_proc_macros::sequencer_latency_histogram;
@@ -33,6 +37,7 @@ use starknet_api::rpc_transaction::{
     RpcDeclareTransaction,
     RpcTransaction,
 };
+use starknet_types_core::felt::Felt;
 use tracing::{debug, info, warn, Span};
 
 use crate::errors::{
@@ -152,32 +157,31 @@ impl Gateway {
             })?;
 
         let mut blocking_task =
-            ProcessTxBlockingTask::new(self, executable_tx, tokio::runtime::Handle::current())
-                .map_err(|e| {
-                    info!(
-                        "Gateway validation failed for tx with signature: {:?} with error: {}",
-                        tx.signature(),
-                        e
-                    );
-                    metric_counters.record_add_tx_failure(&e);
+        ProcessTxBlockingTask::new(self, executable_tx, tokio::runtime::Handle::current())
+            .map_err(|e| {
+                info!(
+                    "Gateway validation failed for tx with signature: {:?} with error: {}",
+                    tx.signature(),
                     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(),
+                );
+                metric_counters.record_add_tx_failure(&e);
                 e
-            );
-            metric_counters.record_add_tx_failure(&e);
-            e
-        })?;
+            })?;
+
+        let account_nonce = blocking_task.get_nonce()?;
+
+        let is_account_tx_in_mempool =
+            if should_query_mempool_for_account_tx(&executable_tx, account_nonce) {
+                self.mempool_client
+                    .account_tx_in_pool_or_recent_block(executable_tx.contract_address())
+                    .await
+            } else {
+                Ok(false)
+            };
 
-        // 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(account_nonce))
+            curr_span.in_scope(|| blocking_task.process_tx(account_nonce, is_account_tx_in_mempool))
         });
         let handle_result = handle.await;
         let nonce = match handle_result {
@@ -253,43 +257,47 @@ impl Gateway {
     }
 }
 
+// Returns true if we should prefetch mempool data to potentially skip stateful validations:
+// applies only to Invoke with nonce == 1 and when the account nonce is zero (pre-deploy state).
+fn should_query_mempool_for_account_tx(
+    executable_tx: &AccountTransaction,
+    account_nonce: Nonce,
+) -> bool {
+    matches!(executable_tx, AccountTransaction::Invoke(_))
+        && executable_tx.nonce() == Nonce(Felt::ONE)
+        && account_nonce == Nonce(Felt::ZERO)
+}
+
 /// 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>,
-    mempool_client: SharedMempoolClient,
     executable_tx: AccountTransaction,
-    runtime: tokio::runtime::Handle,
 }
 
 impl ProcessTxBlockingTask {
-    pub fn new(
-        gateway: &Gateway,
-        executable_tx: AccountTransaction,
-        runtime: tokio::runtime::Handle,
-    ) -> GatewayResult<Self> {
+    pub fn new(gateway: &Gateway, executable_tx: AccountTransaction) -> GatewayResult<Self> {
         let stateful_tx_validator = gateway
             .stateful_tx_validator_factory
             .instantiate_validator(gateway.state_reader_factory.as_ref())?;
-        Ok(Self {
-            stateful_tx_validator,
-            mempool_client: gateway.mempool_client.clone(),
-            executable_tx,
-            runtime,
-        })
+        Ok(Self { stateful_tx_validator, executable_tx })
     }
 
     fn get_nonce(&mut self) -> GatewayResult<Nonce> {
         let address = self.executable_tx.contract_address();
-        self.stateful_tx_validator.get_nonce(address)
+        let nonce = self.stateful_tx_validator.get_nonce(address)?;
+        Ok(nonce)
     }
 
-    fn process_tx(mut self, account_nonce: Nonce) -> GatewayResult<Nonce> {
+    fn process_tx(
+        mut self,
+        account_nonce: Nonce,
+        is_account_tx_in_mempool: Result<bool, MempoolClientError>,
+    ) -> GatewayResult<Nonce> {
         let nonce = self.stateful_tx_validator.extract_state_nonce_and_run_validations(
             &self.executable_tx,
             account_nonce,
-            self.mempool_client,
-            self.runtime,
+            is_account_tx_in_mempool,
         )?;
 
         Ok(nonce)
diff --git a/crates/apollo_gateway/src/gateway_test.rs b/crates/apollo_gateway/src/gateway_test.rs
@@ -118,6 +118,14 @@ fn mock_stateless_transaction_validator() -> MockStatelessTransactionValidatorTr
     mock_stateless_transaction_validator
 }
 
+#[fixture]
+fn default_mock_mempool_client() -> MockMempoolClient {
+    let mut mock_mempool_client = MockMempoolClient::new();
+    // Gateway tests don't require the mempool client to be checked, any value is fine.
+    mock_mempool_client.expect_account_tx_in_pool_or_recent_block().returning(|_| Ok(false));
+    mock_mempool_client
+}
+
 #[fixture]
 fn mock_dependencies() -> MockDependencies {
     let config = GatewayConfig {
@@ -129,7 +137,7 @@ fn mock_dependencies() -> MockDependencies {
     };
     let state_reader_factory =
         local_test_state_reader_factory(CairoVersion::Cairo1(RunnableCairo1::Casm), true);
-    let mock_mempool_client = MockMempoolClient::new();
+    let mock_mempool_client = default_mock_mempool_client();
     let mock_transaction_converter = MockTransactionConverterTrait::new();
     let mock_stateless_transaction_validator = mock_stateless_transaction_validator();
     MockDependencies {
@@ -318,16 +326,12 @@ fn process_tx_task(
     stateful_transaction_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(),
-    )
-    .expect("instantiate_validator should be mocked in tests");
+    let stateful_tx_validator = stateful_transaction_validator_factory
+        .instantiate_validator(state_reader_factory.as_ref())
+        .expect("instantiate_validator should be mocked in tests");
     ProcessTxBlockingTask {
         stateful_tx_validator,
-        mempool_client: Arc::new(MockMempoolClient::new()),
         executable_tx: executable_invoke_tx(invoke_args()),
-        runtime: tokio::runtime::Handle::current(),
     }
 }
 
@@ -572,10 +576,11 @@ async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_f
         .return_once(|_| Ok(Box::new(mock_stateful_transaction_validator)));
 
     let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory);
-    let account_nonce = nonce!(0);
-    let result = tokio::task::spawn_blocking(move || process_tx_task.process_tx(account_nonce))
-        .await
-        .unwrap();
+
+    let result =
+        tokio::task::spawn_blocking(move || process_tx_task.process_tx(nonce!(0), Ok(false)))
+            .await
+            .unwrap();
 
     assert!(result.is_err());
     assert_eq!(result.unwrap_err().code, error_code);
@@ -608,21 +613,22 @@ async fn add_tx_returns_error_when_instantiating_validator_fails(
     mut mock_dependencies: MockDependencies,
     mut mock_stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
 ) {
-    // Prepare transaction conversion to reach instantiation.
     let tx_args = invoke_args();
     setup_transaction_converter_mock(&mut mock_dependencies.mock_transaction_converter, &tx_args);
 
-    // Fail validator instantiation.
     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()
         .return_once(|_| Err(expected_error));
 
-    // Build gateway and inject the failing factory.
     let mut gateway = mock_dependencies.gateway();
     gateway.stateful_tx_validator_factory = Arc::new(mock_stateful_transaction_validator_factory);
 
     let err = gateway.add_tx(tx_args.get_rpc_tx(), p2p_message_metadata()).await.unwrap_err();
     assert_eq!(err.code, error_code);
 }
+
+// With current implementation, validator instantiation occurs inside ProcessTxBlockingTask::new()
+// and unwraps on error. Therefore, add_tx will panic if instantiation fails.
+// Gateway should return an error (not panic) when validator instantiation fails.
diff --git a/crates/apollo_gateway/src/stateful_transaction_validator.rs b/crates/apollo_gateway/src/stateful_transaction_validator.rs
@@ -5,7 +5,7 @@ use apollo_gateway_types::deprecated_gateway_error::{
     StarknetErrorCode,
 };
 use apollo_gateway_types::errors::GatewaySpecError;
-use apollo_mempool_types::communication::SharedMempoolClient;
+use apollo_mempool_types::communication::MempoolClientError;
 use apollo_proc_macros::sequencer_latency_histogram;
 use blockifier::blockifier::stateful_validator::{
     StatefulValidator,
@@ -104,8 +104,12 @@ pub trait StatefulTransactionValidatorTrait: Send {
         &mut self,
         executable_tx: &ExecutableTransaction,
         account_nonce: Nonce,
-        mempool_client: SharedMempoolClient,
-        runtime: tokio::runtime::Handle,
+        is_account_tx_in_mempool: Result<bool, MempoolClientError>,
+    ) -> StatefulTransactionValidatorResult<Nonce>;
+
+    fn get_nonce(
+        &mut self,
+        account_address: ContractAddress,
     ) -> StatefulTransactionValidatorResult<Nonce>;
 
     fn get_nonce(
@@ -126,10 +130,9 @@ impl<B: BlockifierStatefulValidatorTrait + Send> StatefulTransactionValidatorTra
         &mut self,
         executable_tx: &ExecutableTransaction,
         account_nonce: Nonce,
-        mempool_client: SharedMempoolClient,
-        runtime: tokio::runtime::Handle,
+        is_account_tx_in_mempool: Result<bool, MempoolClientError>,
     ) -> StatefulTransactionValidatorResult<Nonce> {
-        self.run_transaction_validations(executable_tx, account_nonce, mempool_client, runtime)?;
+        self.run_transaction_validations(executable_tx, account_nonce, is_account_tx_in_mempool)?;
         Ok(account_nonce)
     }
 
@@ -148,11 +151,10 @@ impl<B: BlockifierStatefulValidatorTrait + Send> StatefulTransactionValidator<B>
         &mut self,
         executable_tx: &ExecutableTransaction,
         account_nonce: Nonce,
-        mempool_client: SharedMempoolClient,
-        runtime: tokio::runtime::Handle,
+        is_account_tx_in_mempool: Result<bool, MempoolClientError>,
     ) -> StatefulTransactionValidatorResult<()> {
         self.validate_state_preconditions(executable_tx, account_nonce)?;
-        self.run_validate_entry_point(executable_tx, account_nonce, mempool_client, runtime)?;
+        self.run_validate_entry_point(executable_tx, account_nonce, is_account_tx_in_mempool)?;
         Ok(())
     }
 
@@ -215,11 +217,10 @@ impl<B: BlockifierStatefulValidatorTrait + Send> StatefulTransactionValidator<B>
         &mut self,
         executable_tx: &ExecutableTransaction,
         account_nonce: Nonce,
-        mempool_client: SharedMempoolClient,
-        runtime: tokio::runtime::Handle,
+        is_account_tx_in_mempool: Result<bool, MempoolClientError>,
     ) -> StatefulTransactionValidatorResult<()> {
         let skip_validate =
-            skip_stateful_validations(executable_tx, account_nonce, mempool_client, runtime)?;
+            skip_stateful_validations(executable_tx, account_nonce, is_account_tx_in_mempool)?;
         let only_query = false;
         let charge_fee = enforce_fee(executable_tx, only_query);
         let strict_nonce_check = false;
@@ -290,8 +291,7 @@ impl<B: BlockifierStatefulValidatorTrait + Send> StatefulTransactionValidator<B>
 fn skip_stateful_validations(
     tx: &ExecutableTransaction,
     account_nonce: Nonce,
-    mempool_client: SharedMempoolClient,
-    runtime: tokio::runtime::Handle,
+    is_account_tx_in_mempool: Result<bool, MempoolClientError>,
 ) -> StatefulTransactionValidatorResult<bool> {
     if let ExecutableTransaction::Invoke(ExecutableInvokeTransaction { tx, .. }) = tx {
         // check if the transaction nonce is 1, meaning it is post deploy_account, and the
@@ -303,8 +303,7 @@ fn skip_stateful_validations(
             // to check if the account exists in the mempool since it means that either it has a
             // deploy_account transaction or transactions with future nonces that passed
             // validations.
-            return runtime
-                .block_on(mempool_client.account_tx_in_pool_or_recent_block(tx.sender_address()))
+            return is_account_tx_in_mempool
                 .map_err(|err| mempool_client_err_to_deprecated_gw_err(&tx.signature(), err))
                 .inspect(|exists| {
                     if *exists {
diff --git a/crates/apollo_gateway/src/stateful_transaction_validator_test.rs b/crates/apollo_gateway/src/stateful_transaction_validator_test.rs
