apollo_gateway: extract async mempool query from blocking task

itamar-starkware · itamar-starkware · commit 14142cc7af81 · 2025-11-12T22:32:40.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::{
@@ -151,32 +156,49 @@ impl Gateway {
                 transaction_converter_err_to_deprecated_gw_err(&tx_signature, e)
             })?;
 
-        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
-                })?;
+        let mut blocking_task = ProcessTxBlockingTask::new(
+            self,
+            executable_tx.clone(),
+            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
+        })?;
 
         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 account_nonce = state_reader
+            .get_account_nonce(account_address)
+            .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 {
@@ -249,12 +271,22 @@ 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,
 }
@@ -268,23 +300,22 @@ impl ProcessTxBlockingTask {
         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,
         })
     }
 
     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,
-            )
-        })?;
+            .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)
     }
 
@@ -307,12 +338,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
@@ -117,6 +117,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 {
@@ -128,7 +136,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 {
@@ -316,11 +324,11 @@ 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(),
     }
@@ -560,12 +568,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();
@@ -601,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_get_state_reader_for_validation()
         .return_once(|_| Box::pin(async { 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 async_trait::async_trait;
 use blockifier::blockifier::stateful_validator::{
@@ -134,8 +134,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(
@@ -156,10 +155,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)
     }
 
@@ -178,11 +176,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(())
     }
 
@@ -253,11 +250,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;
@@ -335,8 +331,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
@@ -348,8 +343,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
