apollo_gateway: extract async mempool query from blocking task

itamar-starkware · itamar-starkware · commit b4115d27af7e · 2025-11-26T14:31:52.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;
@@ -34,6 +38,7 @@ use starknet_api::rpc_transaction::{
     RpcDeclareTransaction,
     RpcTransaction,
 };
+use starknet_types_core::felt::Felt;
 use tracing::{debug, info, warn, Span};
 
 use crate::errors::{
@@ -159,16 +164,14 @@ impl Gateway {
             })?;
 
         let mut blocking_task =
-            ProcessTxBlockingTask::new(self, executable_tx, tokio::runtime::Handle::current())
-                .await
-                .map_err(|e| {
-                    info!(
-                        "Gateway validation failed for tx with signature: {:?} with error: {}",
-                        &tx_signature, e
-                    );
-                    metric_counters.record_add_tx_failure(&e);
-                    e
-                })?;
+            ProcessTxBlockingTask::new(self, executable_tx.clone()).await.map_err(|e| {
+                info!(
+                    "Gateway validation failed for tx with signature: {:?} with error: {}",
+                    &tx_signature, e
+                );
+                metric_counters.record_add_tx_failure(&e);
+                e
+            })?;
 
         let state_reader = blocking_task.get_state_reader().await?;
 
@@ -178,12 +181,26 @@ impl Gateway {
             .await
             .map_err(|e| StarknetError::internal_with_logging("Failed to get account nonce", e))?;
 
+        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)
+            };
+
         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(account_nonce, stateful_tx_validator))
+            curr_span.in_scope(|| {
+                blocking_task.process_tx(
+                    account_nonce,
+                    stateful_tx_validator,
+                    is_account_tx_in_mempool,
+                )
+            })
         });
         let handle_result = handle.await;
         let nonce = match handle_result {
@@ -256,28 +273,31 @@ 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 {
     state_reader_factory: Arc<dyn StateReaderFactory>,
     stateful_tx_validator_factory: Arc<dyn StatefulTransactionValidatorFactoryTrait>,
-    mempool_client: SharedMempoolClient,
     executable_tx: AccountTransaction,
-    runtime: tokio::runtime::Handle,
 }
 
 impl ProcessTxBlockingTask {
-    pub async fn new(
-        gateway: &Gateway,
-        executable_tx: AccountTransaction,
-        runtime: tokio::runtime::Handle,
-    ) -> GatewayResult<Self> {
+    pub async fn new(gateway: &Gateway, executable_tx: AccountTransaction) -> GatewayResult<Self> {
         Ok(Self {
             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,
         })
     }
 
@@ -316,12 +336,12 @@ impl ProcessTxBlockingTask {
         self,
         account_nonce: Nonce,
         mut stateful_tx_validator: Box<dyn StatefulTransactionValidatorTrait>,
+        is_account_tx_in_mempool: Result<bool, MempoolClientError>,
     ) -> GatewayResult<Nonce> {
         let nonce = 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 {
@@ -130,7 +138,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,13 +326,12 @@ async fn run_add_tx_and_extract_metrics(
 async fn process_tx_task(
     stateful_tx_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
 ) -> ProcessTxBlockingTask {
-    let state_reader_factory = Arc::new(MockStateReaderFactory::new());
+    let state_reader_factory: Arc<dyn crate::state_reader::StateReaderFactory> =
+        Arc::new(MockStateReaderFactory::new());
     ProcessTxBlockingTask {
         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(),
     }
 }
 
@@ -562,12 +569,16 @@ 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));
 
     let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory).await;
-    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))
+        process_tx_task.process_tx(
+            nonce!(0),
+            Box::new(mock_stateful_transaction_validator),
+            Ok(false),
+        )
     })
     .await
     .unwrap();
@@ -603,7 +614,6 @@ async fn add_tx_returns_error_when_getting_state_reader_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);
 
@@ -613,7 +623,6 @@ async fn add_tx_returns_error_when_getting_state_reader_fails(
         .expect_get_state_reader_for_validation()
         .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);
 
diff --git a/crates/apollo_gateway/src/stateful_transaction_validator.rs b/crates/apollo_gateway/src/stateful_transaction_validator.rs
@@ -7,7 +7,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 async_trait::async_trait;
 use blockifier::blockifier::stateful_validator::{
@@ -149,8 +149,7 @@ 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(
@@ -171,10 +170,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)
     }
 
@@ -193,11 +191,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(())
     }
 
@@ -294,11 +291,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;
@@ -354,8 +350,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
@@ -367,10 +362,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.
-            // TODO(Itamar): Avoid blocking here by moving this mempool existence check to the async
-            // path prior to spawning the blocking task, and pass the result into validation.
-            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
