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,
@@ -154,7 +155,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 +243,33 @@ 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_transaction_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> {
+        // TODO(Itamar): Extract creating validator to a separate function.
+        let stateful_transaction_validator = gateway
+            .stateful_tx_validator_factory
+            .instantiate_validator(gateway.state_reader_factory.clone())
+            .await?;
+        Ok(Self {
+            stateful_transaction_validator,
             mempool_client: gateway.mempool_client.clone(),
             executable_tx,
             runtime,
-        }
+        })
     }
 
-    fn process_tx(self) -> GatewayResult<Nonce> {
-        // TODO(Itamar): Remove this block_on by extracting the instantiation to the async path.
-        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(
+    fn process_tx(mut self) -> GatewayResult<Nonce> {
+        let nonce = self.stateful_transaction_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,
 };
 use crate::stateless_transaction_validator::MockStatelessTransactionValidatorTrait;
 
@@ -315,12 +316,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,
+    )
+    .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_transaction_validator: stateful_tx_validator,
         mempool_client: Arc::new(MockMempoolClient::new()),
         executable_tx: executable_invoke_tx(invoke_args()),
         runtime: tokio::runtime::Handle::current(),
@@ -567,7 +574,7 @@ 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 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();
 
@@ -598,22 +605,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(|_| 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/stateful_transaction_validator.rs

@@ -143,7 +143,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,
@@ -152,12 +152,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(
@@ -181,7 +181,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,