apollo_gateway: make StatefulTransactionValidatorTrait async

Author: itamar-starkware

crates/apollo_gateway/src/gateway.rs

@@ -32,7 +32,7 @@ use starknet_api::rpc_transaction::{
     RpcDeclareTransaction,
     RpcTransaction,
 };
-use tracing::{debug, warn, Span};
+use tracing::{debug, warn};
 
 use crate::errors::{
     mempool_client_result_to_deprecated_gw_result,
@@ -157,28 +157,10 @@ impl Gateway {
             .await
             .inspect_err(|e| metric_counters.record_add_tx_failure(e))?;
 
-        let curr_span = Span::current();
-        let mempool_client = self.mempool_client.clone();
-        let nonce = tokio::task::spawn_blocking(move || {
-            curr_span.in_scope(|| {
-                stateful_transaction_validator.extract_state_nonce_and_run_validations(
-                    &executable_tx.clone(),
-                    mempool_client,
-                    tokio::runtime::Handle::current(),
-                )
-            })
-        })
-        .await
-        .map_err(|e| {
-            let err = StarknetError {
-                code: StarknetErrorCode::UnknownErrorCode(
-                    "StarknetErrorCode.InternalError".to_string(),
-                ),
-                message: format!("Validation task failed to complete: {e}"),
-            };
-            metric_counters.record_add_tx_failure(&err);
-            err
-        })??;
+        let nonce = stateful_transaction_validator
+            .extract_state_nonce_and_run_validations(&executable_tx, self.mempool_client.clone())
+            .await
+            .inspect_err(|e| metric_counters.record_add_tx_failure(e))?;
 
         let gateway_output = create_gateway_output(&internal_tx);
 

crates/apollo_gateway/src/gateway_test.rs

@@ -557,7 +557,7 @@ async fn add_tx_returns_error_when_run_transaction_validations_fails(
 
     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()

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -11,10 +11,7 @@ use apollo_mempool_types::communication::SharedMempoolClient;
 use apollo_mempool_types::mempool_types::ValidationArgs;
 use apollo_proc_macros::sequencer_latency_histogram;
 use async_trait::async_trait;
-use blockifier::blockifier::stateful_validator::{
-    StatefulValidator,
-    StatefulValidatorTrait as BlockifierStatefulValidatorTrait,
-};
+use blockifier::blockifier::stateful_validator::{StatefulValidator, StatefulValidatorTrait};
 use blockifier::blockifier_versioned_constants::VersionedConstants;
 use blockifier::bouncer::BouncerConfig;
 use blockifier::context::{BlockContext, ChainInfo};
@@ -87,115 +84,96 @@ impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFa
             self.contract_class_manager.clone(),
             Some(GATEWAY_CLASS_CACHE_METRICS),
         );
-
-        let state = CachedState::new(state_reader_and_contract_manager);
-        let mut versioned_constants = VersionedConstants::get_versioned_constants(
-            self.config.versioned_constants_overrides.clone(),
-        );
-        // The validation of a transaction is not affected by the casm hash migration.
-        versioned_constants.enable_casm_hash_migration = false;
-
-        let mut block_info = gateway_fixed_block_state_reader.get_block_info().await?;
-        block_info.block_number = block_info.block_number.unchecked_next();
-        let block_context = BlockContext::new(
-            block_info,
+        Ok(Box::new(StatefulTransactionValidator::new(
+            self.config.clone(),
             self.chain_info.clone(),
-            versioned_constants,
-            BouncerConfig::max(),
-        );
-
-        let blockifier_stateful_tx_validator =
-            BlockifierStatefulValidator::create(state, block_context);
-        Ok(Box::new(StatefulTransactionValidator {
-            config: self.config.clone(),
-            blockifier_stateful_tx_validator,
+            state_reader_and_contract_manager,
             gateway_fixed_block_state_reader,
-        }))
+        )))
     }
 }
 
 #[cfg_attr(test, mockall::automock)]
+#[async_trait]
 pub trait StatefulTransactionValidatorTrait: Send {
-    fn extract_state_nonce_and_run_validations(
+    async fn extract_state_nonce_and_run_validations(
         &mut self,
         executable_tx: &ExecutableTransaction,
         mempool_client: SharedMempoolClient,
-        runtime: tokio::runtime::Handle,
     ) -> StatefulTransactionValidatorResult<Nonce>;
 }
 
-pub struct StatefulTransactionValidator<B: BlockifierStatefulValidatorTrait> {
+pub struct StatefulTransactionValidator {
     config: StatefulTransactionValidatorConfig,
-    blockifier_stateful_tx_validator: B,
+    chain_info: ChainInfo,
+    // Consumed when running the CPU-heavy blockifier validation.
+    // TODO(Itamar): The whole `StatefulTransactionValidator` is never used after
+    // `state_reader_and_contract_manager` is taken. Make it non-optional and discard the
+    // instance after use.
+    state_reader_and_contract_manager:
+        Option<StateReaderAndContractManager<Box<dyn GatewayStateReaderWithCompiledClasses>>>,
     gateway_fixed_block_state_reader: Box<dyn GatewayFixedBlockStateReader>,
 }
 
-impl<B: BlockifierStatefulValidatorTrait + Send> StatefulTransactionValidatorTrait
-    for StatefulTransactionValidator<B>
-{
-    fn extract_state_nonce_and_run_validations(
+#[async_trait]
+impl StatefulTransactionValidatorTrait for StatefulTransactionValidator {
+    async fn extract_state_nonce_and_run_validations(
         &mut self,
         executable_tx: &ExecutableTransaction,
         mempool_client: SharedMempoolClient,
-        runtime: tokio::runtime::Handle,
     ) -> StatefulTransactionValidatorResult<Nonce> {
-        let address = executable_tx.contract_address();
-        let account_nonce = runtime
-            .block_on(self.gateway_fixed_block_state_reader.get_nonce(address))
-            .map_err(|e| {
-                // TODO(noamsp): Fix this. Need to map the errors better.
-                StarknetError::internal_with_signature_logging(
-                    format!("Failed to get nonce for sender address {address}"),
-                    &executable_tx.signature(),
-                    e,
-                )
-            })?;
-        self.run_transaction_validations(executable_tx, account_nonce, mempool_client, runtime)?;
+        let account_nonce = self.extract_nonce(executable_tx).await?;
+        let skip_validate =
+            self.run_pre_validation_checks(executable_tx, account_nonce, mempool_client).await?;
+        self.run_validate_entry_point(executable_tx, skip_validate).await?;
         Ok(account_nonce)
     }
 }
 
-impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidator<B> {
-    fn run_transaction_validations(
+impl StatefulTransactionValidator {
+    fn new(
+        config: StatefulTransactionValidatorConfig,
+        chain_info: ChainInfo,
+        state_reader_and_contract_manager: StateReaderAndContractManager<
+            Box<dyn GatewayStateReaderWithCompiledClasses>,
+        >,
+        gateway_fixed_block_state_reader: Box<dyn GatewayFixedBlockStateReader>,
+    ) -> Self {
+        Self {
+            config,
+            chain_info,
+            state_reader_and_contract_manager: Some(state_reader_and_contract_manager),
+            gateway_fixed_block_state_reader,
+        }
+    }
+
+    fn take_state_reader_and_contract_manager(
         &mut self,
-        executable_tx: &ExecutableTransaction,
-        account_nonce: Nonce,
-        mempool_client: SharedMempoolClient,
-        runtime: tokio::runtime::Handle,
-    ) -> StatefulTransactionValidatorResult<()> {
-        self.validate_state_preconditions(executable_tx, account_nonce)?;
-        runtime.block_on(validate_by_mempool(
-            executable_tx,
-            account_nonce,
-            mempool_client.clone(),
-        ))?;
-        self.run_validate_entry_point(executable_tx, account_nonce, mempool_client, runtime)?;
-        Ok(())
+    ) -> StateReaderAndContractManager<Box<dyn GatewayStateReaderWithCompiledClasses>> {
+        self.state_reader_and_contract_manager.take().expect("Validator was already consumed")
     }
 
-    fn validate_state_preconditions(
+    async fn validate_state_preconditions(
         &self,
         executable_tx: &ExecutableTransaction,
         account_nonce: Nonce,
     ) -> StatefulTransactionValidatorResult<()> {
-        self.validate_resource_bounds(executable_tx)?;
+        self.validate_resource_bounds(executable_tx).await?;
         self.validate_nonce(executable_tx, account_nonce)?;
         Ok(())
     }
 
-    fn validate_resource_bounds(
+    async fn validate_resource_bounds(
         &self,
         executable_tx: &ExecutableTransaction,
     ) -> StatefulTransactionValidatorResult<()> {
         // Skip this validation during the systems bootstrap phase.
         if self.config.validate_resource_bounds {
             // TODO(Arni): getnext_l2_gas_price from the block header.
-            // TODO(Itamar): Replace usage of `blockifier_stateful_tx_validator.block_info()` with
-            // the GW fixed-block provider and then remove `block_info()` from
-            // blockifier::{StatefulValidatorTrait, StatefulValidator}.
             let previous_block_l2_gas_price = self
-                .blockifier_stateful_tx_validator
-                .block_info()
+                .gateway_fixed_block_state_reader
+                .get_block_info()
+                .await?
                 .gas_prices
                 .strk_gas_prices
                 .l2_gas_price;
@@ -265,23 +243,51 @@ impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidator<B> {
     }
 
     #[sequencer_latency_histogram(GATEWAY_VALIDATE_TX_LATENCY, true)]
-    fn run_validate_entry_point(
+    async fn run_validate_entry_point(
         &mut self,
         executable_tx: &ExecutableTransaction,
-        account_nonce: Nonce,
-        mempool_client: SharedMempoolClient,
-        runtime: tokio::runtime::Handle,
+        skip_validate: bool,
     ) -> StatefulTransactionValidatorResult<()> {
-        let skip_validate =
-            skip_stateful_validations(executable_tx, account_nonce, mempool_client, runtime)?;
         let only_query = false;
         let charge_fee = enforce_fee(executable_tx, only_query);
         let strict_nonce_check = false;
         let execution_flags =
             ExecutionFlags { only_query, charge_fee, validate: !skip_validate, strict_nonce_check };
 
         let account_tx = AccountTransaction { tx: executable_tx.clone(), execution_flags };
-        self.blockifier_stateful_tx_validator.validate(account_tx).map_err(|e| StarknetError {
+
+        // Build block context here.
+        let mut versioned_constants = VersionedConstants::get_versioned_constants(
+            self.config.versioned_constants_overrides.clone(),
+        );
+        // The validation of a transaction is not affected by the casm hash migration.
+        versioned_constants.enable_casm_hash_migration = false;
+
+        let mut block_info = self.gateway_fixed_block_state_reader.get_block_info().await?;
+        block_info.block_number = block_info.block_number.unchecked_next();
+        let block_context = BlockContext::new(
+            block_info,
+            self.chain_info.clone(),
+            versioned_constants,
+            BouncerConfig::max(),
+        );
+
+        // Move state into the blocking task and run CPU-heavy validation.
+        let state_reader_and_contract_manager = self.take_state_reader_and_contract_manager();
+
+        tokio::task::spawn_blocking(move || {
+            let state = CachedState::new(state_reader_and_contract_manager);
+            let mut blockifier = BlockifierStatefulValidator::create(state, block_context);
+            blockifier.validate(account_tx)
+        })
+        .await
+        .map_err(|e| StarknetError {
+            code: StarknetErrorCode::UnknownErrorCode(
+                "StarknetErrorCode.InternalError".to_string(),
+            ),
+            message: format!("Blocking task join error: {e}"),
+        })?
+        .map_err(|e| StarknetError {
             code: StarknetErrorCode::KnownErrorCode(KnownStarknetErrorCode::ValidateFailure),
             message: e.to_string(),
         })?;
@@ -321,6 +327,35 @@ impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidator<B> {
         }
         Ok(())
     }
+
+    async fn extract_nonce(
+        &self,
+        executable_tx: &ExecutableTransaction,
+    ) -> StatefulTransactionValidatorResult<Nonce> {
+        let address = executable_tx.contract_address();
+        let account_nonce =
+            self.gateway_fixed_block_state_reader.get_nonce(address).await.map_err(|e| {
+                StarknetError::internal_with_signature_logging(
+                    format!("Failed to get nonce for sender address {address}"),
+                    &executable_tx.signature(),
+                    e,
+                )
+            })?;
+        Ok(account_nonce)
+    }
+
+    async fn run_pre_validation_checks(
+        &self,
+        executable_tx: &ExecutableTransaction,
+        account_nonce: Nonce,
+        mempool_client: SharedMempoolClient,
+    ) -> StatefulTransactionValidatorResult<bool> {
+        self.validate_state_preconditions(executable_tx, account_nonce).await?;
+        validate_by_mempool(executable_tx, account_nonce, mempool_client.clone()).await?;
+        let skip_validate =
+            skip_stateful_validations(executable_tx, account_nonce, mempool_client.clone()).await?;
+        Ok(skip_validate)
+    }
 }
 
 /// Perform transaction validation by the mempool.
@@ -339,11 +374,10 @@ async fn validate_by_mempool(
 /// Check if validation of an invoke transaction should be skipped due to deploy_account not being
 /// processed yet. This feature is used to improve UX for users sending deploy_account + invoke at
 /// once.
-fn skip_stateful_validations(
+async fn skip_stateful_validations(
     tx: &ExecutableTransaction,
     account_nonce: Nonce,
     mempool_client: SharedMempoolClient,
-    runtime: tokio::runtime::Handle,
 ) -> StatefulTransactionValidatorResult<bool> {
     if let ExecutableTransaction::Invoke(ExecutableInvokeTransaction { tx, .. }) = tx {
         // check if the transaction nonce is 1, meaning it is post deploy_account, and the
@@ -355,8 +389,9 @@ 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 mempool_client
+                .account_tx_in_pool_or_recent_block(tx.sender_address())
+                .await
                 .map_err(|err| mempool_client_err_to_deprecated_gw_err(&tx.signature(), err))
                 .inspect(|exists| {
                     if *exists {