apollo_gateway: insert stateful_tx_validator creating to ProcessTxBlockingTask constructor

Author: itamar-starkware

crates/apollo_gateway/src/gateway.rs

@@ -46,6 +46,7 @@ use crate::state_reader::StateReaderFactory;
 use crate::stateful_transaction_validator::{
     StatefulTransactionValidatorFactory,
     StatefulTransactionValidatorFactoryTrait,
+    StatefulTransactionValidatorTrait,
 };
 use crate::stateless_transaction_validator::{
     StatelessTransactionValidator,
@@ -133,13 +134,14 @@ impl Gateway {
         self.stateless_tx_validator.validate(&tx)?;
 
         let tx_signature = tx.signature().clone();
-        let internal_tx =
-            self.transaction_converter.convert_rpc_tx_to_internal_rpc_tx(tx).await.map_err(
-                |e| {
-                    warn!("Failed to convert RPC transaction to internal RPC transaction: {}", e);
-                    transaction_converter_err_to_deprecated_gw_err(&tx_signature, e)
-                },
-            )?;
+        let internal_tx = self
+            .transaction_converter
+            .convert_rpc_tx_to_internal_rpc_tx(tx.clone())
+            .await
+            .map_err(|e| {
+                warn!("Failed to convert RPC transaction to internal RPC transaction: {}", e);
+                transaction_converter_err_to_deprecated_gw_err(&tx_signature, e)
+            })?;
 
         let executable_tx = self
             .transaction_converter
@@ -154,7 +156,16 @@ impl Gateway {
             })?;
 
         let blocking_task =
-            ProcessTxBlockingTask::new(self, executable_tx, tokio::runtime::Handle::current());
+            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
+                })?;
         // Run the blocking task in the current span.
         let curr_span = Span::current();
         let handle =
@@ -233,35 +244,32 @@ 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>,
+    stateful_tx_validator: Box<dyn StatefulTransactionValidatorTrait + Send>,
     mempool_client: SharedMempoolClient,
     executable_tx: AccountTransaction,
     runtime: tokio::runtime::Handle,
 }
 
 impl ProcessTxBlockingTask {
-    pub fn new(
+    pub async 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(),
+    ) -> GatewayResult<Self> {
+        let stateful_tx_validator = gateway
+            .stateful_tx_validator_factory
+            .instantiate_validator(gateway.state_reader_factory.as_ref())
+            .await?;
+        Ok(Self {
+            stateful_tx_validator,
             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.as_ref()),
-        )?;
-
-        let nonce = stateful_transaction_validator.extract_state_nonce_and_run_validations(
+    fn process_tx(mut self) -> GatewayResult<Nonce> {
+        let nonce = self.stateful_tx_validator.extract_state_nonce_and_run_validations(
             &self.executable_tx,
             self.mempool_client,
             self.runtime,

crates/apollo_gateway/src/gateway_test.rs

@@ -98,6 +98,7 @@ use crate::state_reader_test_utils::{local_test_state_reader_factory, TestStateR
 use crate::stateful_transaction_validator::{
     MockStatefulTransactionValidatorFactoryTrait,
     MockStatefulTransactionValidatorTrait,
+    StatefulTransactionValidatorFactoryTrait,
     StatefulTransactionValidatorTrait,
 };
 use crate::stateless_transaction_validator::MockStatelessTransactionValidatorTrait;
@@ -316,12 +317,18 @@ async fn run_add_tx_and_extract_metrics(
     AddTxResults { result, metric_handle_for_queries, metrics }
 }
 
-fn process_tx_task(
+async fn process_tx_task(
     stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
 ) -> ProcessTxBlockingTask {
+    let state_reader_factory = Arc::new(MockStateReaderFactory::new());
+    let stateful_tx_validator = StatefulTransactionValidatorFactoryTrait::instantiate_validator(
+        &stateful_transaction_validator_factory,
+        state_reader_factory.as_ref(),
+    )
+    .await
+    .expect("instantiate_validator should be mocked in tests");
     ProcessTxBlockingTask {
-        stateful_tx_validator_factory: Arc::new(stateful_transaction_validator_factory),
-        state_reader_factory: Arc::new(MockStateReaderFactory::new()),
+        stateful_tx_validator,
         mempool_client: Arc::new(MockMempoolClient::new()),
         executable_tx: executable_invoke_tx(invoke_args()),
         runtime: tokio::runtime::Handle::current(),
@@ -572,7 +579,7 @@ async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_f
         })
     });
 
-    let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory);
+    let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory).await;
 
     let result = tokio::task::spawn_blocking(move || process_tx_task.process_tx()).await.unwrap();
 
@@ -603,22 +610,25 @@ 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_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(), // Message is not checked
-    };
+    let expected_error = StarknetError { code: error_code.clone(), message: "placeholder".into() };
     mock_stateful_transaction_validator_factory
         .expect_instantiate_validator()
         .return_once(|_| Box::pin(async { Err::<_, _>(expected_error) }));
 
-    let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory);
-
-    let result = tokio::task::spawn_blocking(move || process_tx_task.process_tx()).await.unwrap();
+    // 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);
 
-    assert!(result.is_err());
-    assert_eq!(result.unwrap_err().code, error_code);
+    let err = gateway.add_tx(tx_args.get_rpc_tx(), p2p_message_metadata()).await.unwrap_err();
+    assert_eq!(err.code, error_code);
 }

crates/apollo_gateway/src/state_reader.rs

@@ -19,6 +19,7 @@ pub trait MempoolStateReader: BlockifierStateReader + Send + Sync {
     async fn get_block_info(&self) -> StateResult<BlockInfo>;
 }
 
+#[async_trait]
 #[cfg_attr(test, automock)]
 #[async_trait]
 pub trait StateReaderFactory: Send + Sync {

crates/apollo_gateway/src/state_reader_test_utils.rs

@@ -1,5 +1,5 @@
 use apollo_state_sync_types::communication::StateSyncClientResult;
-use axum::async_trait;
+use async_trait::async_trait;
 use blockifier::context::BlockContext;
 use blockifier::execution::contract_class::RunnableCompiledClass;
 use blockifier::state::errors::StateError;

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -141,7 +141,7 @@ impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFa
 }
 
 #[cfg_attr(test, mockall::automock)]
-pub trait StatefulTransactionValidatorTrait {
+pub trait StatefulTransactionValidatorTrait: Send {
     fn extract_state_nonce_and_run_validations(
         &mut self,
         executable_tx: &ExecutableTransaction,
@@ -150,12 +150,12 @@ pub trait StatefulTransactionValidatorTrait {
     ) -> StatefulTransactionValidatorResult<Nonce>;
 }
 
-pub struct StatefulTransactionValidator<B: BlockifierStatefulValidatorTrait> {
+pub struct StatefulTransactionValidator<B: BlockifierStatefulValidatorTrait + Send> {
     config: StatefulTransactionValidatorConfig,
     blockifier_stateful_tx_validator: B,
 }
 
-impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidatorTrait
+impl<B: BlockifierStatefulValidatorTrait + Send> StatefulTransactionValidatorTrait
     for StatefulTransactionValidator<B>
 {
     fn extract_state_nonce_and_run_validations(
@@ -179,7 +179,7 @@ impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidatorTrait
     }
 }
 
-impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidator<B> {
+impl<B: BlockifierStatefulValidatorTrait + Send> StatefulTransactionValidator<B> {
     fn run_transaction_validations(
         &mut self,
         executable_tx: &ExecutableTransaction,