apollo_gateway: extract async mempool query from blocking task

Author: itamar-starkware

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;
@@ -150,8 +154,15 @@ impl Gateway {
                 transaction_converter_err_to_deprecated_gw_err(&tx_signature, e)
             })?;
 
+        // Prefetch mempool result in async context to avoid blocking-on-async in a blocking
+        // thread.
+        let is_account_tx_in_mempool = self
+            .mempool_client
+            .account_tx_in_pool_or_recent_block(executable_tx.contract_address())
+            .await;
+
         let blocking_task =
-            ProcessTxBlockingTask::new(self, executable_tx, tokio::runtime::Handle::current());
+            ProcessTxBlockingTask::new(self, executable_tx, is_account_tx_in_mempool);
         // Run the blocking task in the current span.
         let curr_span = Span::current();
         let handle =
@@ -235,23 +246,21 @@ impl Gateway {
 struct ProcessTxBlockingTask {
     stateful_tx_validator_factory: Arc<dyn StatefulTransactionValidatorFactoryTrait>,
     state_reader_factory: Arc<dyn StateReaderFactory>,
-    mempool_client: SharedMempoolClient,
     executable_tx: AccountTransaction,
-    runtime: tokio::runtime::Handle,
+    is_account_tx_in_mempool: Result<bool, MempoolClientError>,
 }
 
 impl ProcessTxBlockingTask {
     pub fn new(
         gateway: &Gateway,
         executable_tx: AccountTransaction,
-        runtime: tokio::runtime::Handle,
+        is_account_tx_in_mempool: Result<bool, MempoolClientError>,
     ) -> Self {
         Self {
             stateful_tx_validator_factory: gateway.stateful_tx_validator_factory.clone(),
             state_reader_factory: gateway.state_reader_factory.clone(),
-            mempool_client: gateway.mempool_client.clone(),
             executable_tx,
-            runtime,
+            is_account_tx_in_mempool,
         }
     }
 
@@ -262,8 +271,7 @@ impl ProcessTxBlockingTask {
 
         let nonce = stateful_transaction_validator.extract_state_nonce_and_run_validations(
             &self.executable_tx,
-            self.mempool_client,
-            self.runtime,
+            self.is_account_tx_in_mempool,
         )?;
 
         Ok(nonce)

crates/apollo_gateway/src/gateway_test.rs

@@ -117,6 +117,15 @@ fn mock_stateless_transaction_validator() -> MockStatelessTransactionValidatorTr
     mock_stateless_transaction_validator
 }
 
+#[fixture]
+fn default_mock_mempool_client() -> MockMempoolClient {
+    let mut mock = MockMempoolClient::new();
+    // TODO(itamars): Add test which mempool is checked.
+    mock.expect_account_tx_in_pool_or_recent_block()
+        .returning(|_| Err(MempoolClientError::MempoolError(MempoolError::MempoolFull)));
+    mock
+}
+
 #[fixture]
 fn mock_dependencies() -> MockDependencies {
     let config = GatewayConfig {
@@ -128,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 {
@@ -315,13 +324,13 @@ async fn run_add_tx_and_extract_metrics(
 
 fn process_tx_task(
     stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
+    is_account_tx_in_mempool: Result<bool, MempoolClientError>,
 ) -> ProcessTxBlockingTask {
     ProcessTxBlockingTask {
         stateful_tx_validator_factory: Arc::new(stateful_transaction_validator_factory),
         state_reader_factory: Arc::new(MockStateReaderFactory::new()),
-        mempool_client: Arc::new(MockMempoolClient::new()),
         executable_tx: executable_invoke_tx(invoke_args()),
-        runtime: tokio::runtime::Handle::current(),
+        is_account_tx_in_mempool,
     }
 }
 
@@ -559,13 +568,13 @@ 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(|_| Ok(Box::new(mock_stateful_transaction_validator)));
 
-    let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory);
+    let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory, Ok(false));
 
     let result = tokio::task::spawn_blocking(move || process_tx_task.process_tx()).await.unwrap();
 
@@ -608,7 +617,7 @@ async fn process_tx_returns_error_when_instantiating_validator_fails(
         .expect_instantiate_validator()
         .return_once(|_| Err(expected_error));
 
-    let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory);
+    let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory, Ok(false));
 
     let result = tokio::task::spawn_blocking(move || process_tx_task.process_tx()).await.unwrap();
 

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,
@@ -103,8 +103,7 @@ pub trait StatefulTransactionValidatorTrait {
     fn extract_state_nonce_and_run_validations(
         &mut self,
         executable_tx: &ExecutableTransaction,
-        mempool_client: SharedMempoolClient,
-        runtime: tokio::runtime::Handle,
+        is_account_tx_in_mempool: Result<bool, MempoolClientError>,
     ) -> StatefulTransactionValidatorResult<Nonce>;
 }
 
@@ -119,8 +118,7 @@ impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidatorTrait
     fn extract_state_nonce_and_run_validations(
         &mut self,
         executable_tx: &ExecutableTransaction,
-        mempool_client: SharedMempoolClient,
-        runtime: tokio::runtime::Handle,
+        is_account_tx_in_mempool: Result<bool, MempoolClientError>,
     ) -> StatefulTransactionValidatorResult<Nonce> {
         let address = executable_tx.contract_address();
         let account_nonce =
@@ -132,7 +130,7 @@ impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidatorTrait
                     e,
                 )
             })?;
-        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)
     }
 }
@@ -142,11 +140,10 @@ impl<B: BlockifierStatefulValidatorTrait> 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(())
     }
 
@@ -209,11 +206,10 @@ impl<B: BlockifierStatefulValidatorTrait> 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;
@@ -284,8 +280,7 @@ impl<B: BlockifierStatefulValidatorTrait> 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
@@ -297,8 +292,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 {

crates/apollo_gateway/src/stateful_transaction_validator_test.rs

@@ -1,12 +1,9 @@
-use std::sync::Arc;
-
 use apollo_gateway_config::config::StatefulTransactionValidatorConfig;
 use apollo_gateway_types::deprecated_gateway_error::{
     KnownStarknetErrorCode,
     StarknetError,
     StarknetErrorCode,
 };
-use apollo_mempool_types::communication::MockMempoolClient;
 use blockifier::blockifier::stateful_validator::{
     MockStatefulValidatorTrait as MockBlockifierStatefulValidatorTrait,
     StatefulValidatorError as BlockifierStatefulValidatorError,
@@ -59,26 +56,17 @@ async fn test_get_nonce_fail_on_extract_state_nonce_and_run_validations() {
         )))
     });
 
-    let mut mock_mempool_client = MockMempoolClient::new();
-    mock_mempool_client.expect_account_tx_in_pool_or_recent_block().returning(|_| {
-        // 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();
-
     let mut stateful_validator = StatefulTransactionValidator {
         config: StatefulTransactionValidatorConfig::default(),
         blockifier_stateful_tx_validator: mock_blockifier_validator,
     };
 
+    // The test fails before the mempool is checked, so we fix it to Ok(false).
+    let is_account_tx_in_mempool = Ok(false);
     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,
-        )
+        stateful_validator
+            .extract_state_nonce_and_run_validations(&executable_tx, is_account_tx_in_mempool)
     })
     .await
     .unwrap();
@@ -131,25 +119,16 @@ async fn test_extract_state_nonce_and_run_validations(
     mock_blockifier_validator.expect_validate().return_once(|_| expected_result.map(|_| ()));
     mock_blockifier_validator.expect_block_info().return_const(BlockInfo::default());
 
-    let mut mock_mempool_client = MockMempoolClient::new();
-    mock_mempool_client.expect_account_tx_in_pool_or_recent_block().returning(|_| {
-        // 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();
-
     let mut stateful_validator = StatefulTransactionValidator {
         config: StatefulTransactionValidatorConfig::default(),
         blockifier_stateful_tx_validator: mock_blockifier_validator,
     };
 
+    // Mempool is checked only when tx nonce is 1; Not the case here, so set a fixed value.
+    let is_account_tx_in_mempool = Ok(false);
     let result = tokio::task::spawn_blocking(move || {
-        stateful_validator.extract_state_nonce_and_run_validations(
-            &executable_tx,
-            mempool_client,
-            runtime,
-        )
+        stateful_validator
+            .extract_state_nonce_and_run_validations(&executable_tx, is_account_tx_in_mempool)
     })
     .await
     .unwrap();
@@ -232,25 +211,17 @@ async fn test_skip_validate(
         .returning(|_| Ok(()));
     mock_blockifier_validator.expect_block_info().return_const(BlockInfo::default());
 
-    let mut mock_mempool_client = MockMempoolClient::new();
-    mock_mempool_client
-        .expect_account_tx_in_pool_or_recent_block()
-        .returning(move |_| Ok(contains_tx));
-    let mempool_client = Arc::new(mock_mempool_client);
-
-    let runtime = tokio::runtime::Handle::current();
-
     let mut stateful_validator = StatefulTransactionValidator {
         config: StatefulTransactionValidatorConfig::default(),
         blockifier_stateful_tx_validator: mock_blockifier_validator,
     };
 
+    let is_account_tx_in_mempool = Ok(contains_tx);
     tokio::task::spawn_blocking(move || {
         let _ = stateful_validator.run_transaction_validations(
             &executable_tx,
             sender_nonce,
-            mempool_client,
-            runtime,
+            is_account_tx_in_mempool,
         );
     })
     .await
@@ -352,12 +323,13 @@ async fn validate_resource_bounds(
         blockifier_stateful_tx_validator: mock_blockifier_validator,
     };
 
+    // Presence is irrelevant for resource-bound validation; fix to Ok(false) for determinism.
+    let is_account_tx_in_mempool = Ok(false);
     let result = tokio::task::spawn_blocking(move || {
         stateful_validator.run_transaction_validations(
             &executable_tx,
             account_nonce,
-            Arc::new(MockMempoolClient::new()),
-            tokio::runtime::Handle::current(),
+            is_account_tx_in_mempool,
         )
     })
     .await
@@ -401,12 +373,14 @@ async fn test_is_valid_nonce(
         nonce: tx_nonce,
         resource_bounds: ValidResourceBounds::create_for_testing(),
     ));
+
+    // Mempool is checked only when account nonce is 0; Not the case here, so set a fixed value.
+    let is_account_tx_in_mempool = Ok(false);
     let result = tokio::task::spawn_blocking(move || {
         stateful_validator.run_transaction_validations(
             &executable_tx,
             account_nonce,
-            Arc::new(MockMempoolClient::new()),
-            tokio::runtime::Handle::current(),
+            is_account_tx_in_mempool,
         )
     })
     .await
@@ -449,12 +423,13 @@ async fn test_reject_future_declares(
         blockifier_stateful_tx_validator: mock_blockifier_validator,
     };
 
+    // Mempool is checked only when account nonce is 0; Not the case here, so set a fixed value.
+    let is_account_tx_in_mempool = Ok(false);
     let result = tokio::task::spawn_blocking(move || {
         stateful_validator.run_transaction_validations(
             &executable_tx,
             nonce!(account_nonce),
-            Arc::new(MockMempoolClient::new()),
-            tokio::runtime::Handle::current(),
+            is_account_tx_in_mempool,
         )
     })
     .await