apollo_gateway: remove processing tx block

Author: itamar-starkware

crates/apollo_gateway/src/gateway.rs

@@ -26,15 +26,13 @@ use apollo_proc_macros::sequencer_latency_histogram;
 use apollo_state_sync_types::communication::SharedStateSyncClient;
 use axum::async_trait;
 use blockifier::state::contract_class_manager::ContractClassManager;
-use starknet_api::core::Nonce;
-use starknet_api::executable_transaction::AccountTransaction;
 use starknet_api::rpc_transaction::{
     InternalRpcTransaction,
     InternalRpcTransactionWithoutTxHash,
     RpcDeclareTransaction,
     RpcTransaction,
 };
-use tracing::{debug, info, warn, Span};
+use tracing::{debug, warn, Span};
 
 use crate::errors::{
     mempool_client_result_to_deprecated_gw_result,
@@ -153,33 +151,34 @@ impl Gateway {
                 transaction_converter_err_to_deprecated_gw_err(&tx_signature, e)
             })?;
 
-        let blocking_task =
-            ProcessTxBlockingTask::new(self, executable_tx, tokio::runtime::Handle::current());
-        // Run the blocking task in the current span.
+        let mut stateful_transaction_validator = self
+            .stateful_tx_validator_factory
+            .instantiate_validator(self.state_reader_factory.clone())
+            .await
+            .inspect_err(|e| metric_counters.record_add_tx_failure(e))?;
+
         let curr_span = Span::current();
-        let handle =
-            tokio::task::spawn_blocking(move || curr_span.in_scope(|| blocking_task.process_tx()));
-        let handle_result = handle.await;
-        let nonce = match handle_result {
-            Ok(Ok(nonce)) => nonce,
-            Ok(Err(starknet_err)) => {
-                info!(
-                    "Gateway validation failed for tx with signature: {:?} with error: {}",
-                    &tx_signature, starknet_err
-                );
-                metric_counters.record_add_tx_failure(&starknet_err);
-                return Err(starknet_err);
-            }
-            Err(join_err) => {
-                let err = StarknetError::internal_with_signature_logging(
-                    "Failed to process tx",
-                    &tx_signature,
-                    join_err,
-                );
-                metric_counters.record_add_tx_failure(&err);
-                return Err(err);
-            }
-        };
+        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 gateway_output = create_gateway_output(&internal_tx);
 
@@ -230,46 +229,6 @@ impl Gateway {
     }
 }
 
-/// CPU-intensive transaction processing, spawned in a blocking thread to avoid blocking other tasks
-/// from running.
-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,
-}
-
-impl ProcessTxBlockingTask {
-    pub fn new(
-        gateway: &Gateway,
-        executable_tx: AccountTransaction,
-        runtime: tokio::runtime::Handle,
-    ) -> 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,
-        }
-    }
-
-    fn process_tx(self) -> GatewayResult<Nonce> {
-        let mut stateful_transaction_validator = self.runtime.block_on(
-            self.stateful_tx_validator_factory.instantiate_validator(self.state_reader_factory),
-        )?;
-
-        let nonce = stateful_transaction_validator.extract_state_nonce_and_run_validations(
-            &self.executable_tx,
-            self.mempool_client,
-            self.runtime,
-        )?;
-
-        Ok(nonce)
-    }
-}
-
 pub fn create_gateway(
     config: GatewayConfig,
     shared_state_sync_client: SharedStateSyncClient,

crates/apollo_gateway/src/gateway_fixed_block_state_reader.rs

@@ -1,9 +1,15 @@
 use std::sync::Arc;
 
 use apollo_gateway_types::deprecated_gateway_error::StarknetError;
-use apollo_state_sync_types::communication::{SharedStateSyncClient, StateSyncClient};
+use apollo_state_sync_types::communication::{
+    SharedStateSyncClient,
+    StateSyncClient,
+    StateSyncClientError,
+};
+use apollo_state_sync_types::errors::StateSyncError;
 use async_trait::async_trait;
 use starknet_api::block::{BlockInfo, BlockNumber, GasPriceVector, GasPrices};
+use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::data_availability::L1DataAvailabilityMode;
 use tokio::sync::OnceCell;
 
@@ -14,6 +20,7 @@ pub type StarknetResult<T> = Result<T, StarknetError>;
 #[async_trait]
 pub trait GatewayFixedBlockStateReader: Send + Sync {
     async fn get_block_info(&self) -> StarknetResult<BlockInfo>;
+    async fn get_nonce(&self, contract_address: ContractAddress) -> StarknetResult<Nonce>;
 }
 
 pub struct GatewayFixedBlockSyncStateClient {
@@ -67,4 +74,14 @@ impl GatewayFixedBlockStateReader for GatewayFixedBlockSyncStateClient {
             .await
             .cloned()
     }
+
+    async fn get_nonce(&self, contract_address: ContractAddress) -> StarknetResult<Nonce> {
+        match self.state_sync_client.get_nonce_at(self.block_number, contract_address).await {
+            Ok(nonce) => Ok(nonce),
+            Err(StateSyncClientError::StateSyncError(StateSyncError::ContractNotFound(_))) => {
+                Ok(Nonce::default())
+            }
+            Err(e) => Err(StarknetError::internal_with_logging("Failed to get nonce", e)),
+        }
+    }
 }

crates/apollo_gateway/src/gateway_test.rs

@@ -81,7 +81,7 @@ use strum::VariantNames;
 use tempfile::TempDir;
 
 use crate::errors::{GatewayResult, StatelessTransactionValidatorError};
-use crate::gateway::{Gateway, ProcessTxBlockingTask};
+use crate::gateway::Gateway;
 use crate::metrics::{
     register_metrics,
     GatewayMetricHandle,
@@ -322,18 +322,6 @@ async fn run_add_tx_and_extract_metrics(
     AddTxResults { result, metric_handle_for_queries, metrics }
 }
 
-fn process_tx_task(
-    stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
-) -> 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(),
-    }
-}
-
 // Gateway spec errors tests.
 // TODO(Arni): Add tests for all the error cases. Check the response (use `into_response` on the
 // result of `add_tx`).
@@ -556,10 +544,11 @@ fn test_full_cycle_dump_deserialize_authorized_declarer_accounts(
     StarknetErrorCode::UnknownErrorCode("StarknetErrorCode.InternalError".into())
 )]
 #[tokio::test]
-async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_fails(
+async fn add_tx_returns_error_when_run_transaction_validations_fails(
     #[case] error_code: StarknetErrorCode,
     mut mock_stateful_transaction_validator: MockStatefulTransactionValidatorTrait,
     mut mock_stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
+    mut mock_dependencies: MockDependencies,
 ) {
     let expected_error = StarknetError {
         code: error_code.clone(),
@@ -574,9 +563,18 @@ async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_f
         .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 tx_args = invoke_args();
+    setup_transaction_converter_mock(&mut mock_dependencies.mock_transaction_converter, &tx_args);
+    let gateway = Gateway {
+        config: Arc::new(mock_dependencies.config),
+        stateless_tx_validator: Arc::new(mock_dependencies.mock_stateless_transaction_validator),
+        stateful_tx_validator_factory: Arc::new(mock_stateful_transaction_validator_factory),
+        state_reader_factory: Arc::new(MockStateReaderFactory::new()),
+        mempool_client: Arc::new(mock_dependencies.mock_mempool_client),
+        transaction_converter: Arc::new(mock_dependencies.mock_transaction_converter),
+    };
 
-    let result = tokio::task::spawn_blocking(move || process_tx_task.process_tx()).await.unwrap();
+    let result = gateway.add_tx(tx_args.get_rpc_tx(), None).await;
 
     assert!(result.is_err());
     assert_eq!(result.unwrap_err().code, error_code);
@@ -605,8 +603,9 @@ async fn stateless_transaction_validator_error(mut mock_dependencies: MockDepend
 
 #[rstest]
 #[tokio::test]
-async fn process_tx_returns_error_when_instantiating_validator_fails(
+async fn add_tx_returns_error_when_instantiating_validator_fails(
     mut mock_stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
+    mut mock_dependencies: MockDependencies,
 ) {
     let error_code = StarknetErrorCode::UnknownErrorCode("StarknetErrorCode.InternalError".into());
     let expected_error = StarknetError {
@@ -617,9 +616,18 @@ 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 tx_args = invoke_args();
+    setup_transaction_converter_mock(&mut mock_dependencies.mock_transaction_converter, &tx_args);
+    let gateway = Gateway {
+        config: Arc::new(mock_dependencies.config),
+        stateless_tx_validator: Arc::new(mock_dependencies.mock_stateless_transaction_validator),
+        stateful_tx_validator_factory: Arc::new(mock_stateful_transaction_validator_factory),
+        state_reader_factory: Arc::new(MockStateReaderFactory::new()),
+        mempool_client: Arc::new(mock_dependencies.mock_mempool_client),
+        transaction_converter: Arc::new(mock_dependencies.mock_transaction_converter),
+    };
 
-    let result = tokio::task::spawn_blocking(move || process_tx_task.process_tx()).await.unwrap();
+    let result = gateway.add_tx(tx_args.get_rpc_tx(), None).await;
 
     assert!(result.is_err());
     assert_eq!(result.unwrap_err().code, error_code);

crates/apollo_gateway/src/rpc_state_reader.rs

@@ -210,6 +210,27 @@ impl GatewayFixedBlockStateReader for RpcStateReader {
         })?;
         Ok(block_info)
     }
+
+    async fn get_nonce(&self, contract_address: ContractAddress) -> StarknetResult<Nonce> {
+        let get_nonce_params = GetNonceParams { block_id: self.block_id, contract_address };
+        let reader = self.clone();
+        let result = tokio::task::spawn_blocking(move || {
+            reader.send_rpc_request("starknet_getNonce", get_nonce_params)
+        })
+        .await
+        .map_err(|e| StarknetError::internal_with_logging("JoinError", e))?;
+
+        match result {
+            Ok(value) => {
+                let nonce: Nonce = serde_json::from_value(value).map_err(|e| {
+                    StarknetError::internal_with_logging("Failed to parse nonce", e)
+                })?;
+                Ok(nonce)
+            }
+            Err(RPCStateReaderError::ContractAddressNotFound(_)) => Ok(Nonce::default()),
+            Err(e) => Err(StarknetError::internal_with_logging("RPC error", e)),
+        }
+    }
 }
 
 #[async_trait]

crates/apollo_gateway/src/state_reader_test_utils.rs

@@ -79,6 +79,10 @@ impl GatewayFixedBlockStateReader for TestGatewayFixedBlockStateReader {
     async fn get_block_info(&self) -> StarknetResult<BlockInfo> {
         Ok(self.block_info.clone())
     }
+
+    async fn get_nonce(&self, _contract_address: ContractAddress) -> StarknetResult<Nonce> {
+        Ok(Nonce::default())
+    }
 }
 
 pub struct TestStateReaderFactory {

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -35,6 +35,7 @@ use starknet_types_core::felt::Felt;
 use tracing::debug;
 
 use crate::errors::{mempool_client_err_to_deprecated_gw_err, StatefulTransactionValidatorResult};
+use crate::gateway_fixed_block_state_reader::GatewayFixedBlockStateReader;
 use crate::metrics::{GATEWAY_CLASS_CACHE_METRICS, GATEWAY_VALIDATE_TX_LATENCY};
 use crate::state_reader::{GatewayStateReaderWithCompiledClasses, StateReaderFactory};
 
@@ -108,12 +109,13 @@ impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFa
         Ok(Box::new(StatefulTransactionValidator {
             config: self.config.clone(),
             blockifier_stateful_tx_validator,
+            gateway_fixed_block_state_reader,
         }))
     }
 }
 
 #[cfg_attr(test, mockall::automock)]
-pub trait StatefulTransactionValidatorTrait {
+pub trait StatefulTransactionValidatorTrait: Send {
     fn extract_state_nonce_and_run_validations(
         &mut self,
         executable_tx: &ExecutableTransaction,
@@ -125,9 +127,10 @@ pub trait StatefulTransactionValidatorTrait {
 pub struct StatefulTransactionValidator<B: BlockifierStatefulValidatorTrait> {
     config: StatefulTransactionValidatorConfig,
     blockifier_stateful_tx_validator: B,
+    gateway_fixed_block_state_reader: Box<dyn GatewayFixedBlockStateReader>,
 }
 
-impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidatorTrait
+impl<B: BlockifierStatefulValidatorTrait + Send> StatefulTransactionValidatorTrait
     for StatefulTransactionValidator<B>
 {
     fn extract_state_nonce_and_run_validations(
@@ -137,22 +140,23 @@ impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidatorTrait
         runtime: tokio::runtime::Handle,
     ) -> StatefulTransactionValidatorResult<Nonce> {
         let address = executable_tx.contract_address();
-        let account_nonce =
-            self.blockifier_stateful_tx_validator.get_nonce(address).map_err(|e| {
+        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)?;
+        self.run_transaction_validation(executable_tx, account_nonce, mempool_client, runtime)?;
         Ok(account_nonce)
     }
 }
 
 impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidator<B> {
-    fn run_transaction_validations(
+    fn run_transaction_validation(
         &mut self,
         executable_tx: &ExecutableTransaction,
         account_nonce: Nonce,