From 2a1b6aa69443449b28d3bd54beb12e744f40cd51 Mon Sep 17 00:00:00 2001
From: Arni Hod <arnon@starkware.co>
Date: Wed, 25 Dec 2024 13:55:34 +0200
Subject: [PATCH] refactor(starknet_gateway): gateway business logic towards
 bench test

---
 crates/starknet_gateway/src/gateway.rs        | 58 ++++++++++++-------
 .../stateful_transaction_validator_test.rs    |  2 +-
 2 files changed, 37 insertions(+), 23 deletions(-)

diff --git a/crates/starknet_gateway/src/gateway.rs b/crates/starknet_gateway/src/gateway.rs
index 102e2cf0b6d..f7577aa036a 100644
--- a/crates/starknet_gateway/src/gateway.rs
+++ b/crates/starknet_gateway/src/gateway.rs
@@ -26,25 +26,21 @@ use crate::utils::compile_contract_and_build_executable_tx;
 #[path = "gateway_test.rs"]
 pub mod gateway_test;
 
-pub struct Gateway {
-    pub config: GatewayConfig,
+struct GatewayBusinessLogic {
     pub stateless_tx_validator: Arc<StatelessTransactionValidator>,
     pub stateful_tx_validator: Arc<StatefulTransactionValidator>,
     pub state_reader_factory: Arc<dyn StateReaderFactory>,
     pub gateway_compiler: Arc<GatewayCompiler>,
-    pub mempool_client: SharedMempoolClient,
     pub chain_info: ChainInfo,
 }
 
-impl Gateway {
+impl GatewayBusinessLogic {
     pub fn new(
         config: GatewayConfig,
         state_reader_factory: Arc<dyn StateReaderFactory>,
         gateway_compiler: GatewayCompiler,
-        mempool_client: SharedMempoolClient,
     ) -> Self {
         Self {
-            config: config.clone(),
             stateless_tx_validator: Arc::new(StatelessTransactionValidator {
                 config: config.stateless_tx_validator_config.clone(),
             }),
@@ -53,31 +49,49 @@ impl Gateway {
             }),
             state_reader_factory,
             gateway_compiler: Arc::new(gateway_compiler),
-            mempool_client,
             chain_info: config.chain_info.clone(),
         }
     }
 
+    pub async fn add_tx(&self, tx: RpcTransaction) -> GatewayResult<AddTransactionArgs> {
+        info!("Processing tx");
+        let blocking_task = ProcessTxBlockingTask::new(self, tx);
+        // Run the blocking task in the current span.
+        let curr_span = Span::current();
+        tokio::task::spawn_blocking(move || curr_span.in_scope(|| blocking_task.process_tx()))
+            .await
+            .map_err(|join_err| {
+                error!("Failed to process tx: {}", join_err);
+                GatewaySpecError::UnexpectedError { data: "Internal server error".to_owned() }
+            })?
+    }
+}
+
+pub struct Gateway {
+    pub mempool_client: SharedMempoolClient,
+    business_logic: GatewayBusinessLogic,
+}
+
+impl Gateway {
+    pub fn new(
+        config: GatewayConfig,
+        state_reader_factory: Arc<dyn StateReaderFactory>,
+        gateway_compiler: GatewayCompiler,
+        mempool_client: SharedMempoolClient,
+    ) -> Self {
+        let business_logic =
+            GatewayBusinessLogic::new(config, state_reader_factory, gateway_compiler);
+        Self { business_logic, mempool_client }
+    }
+
     #[instrument(skip(self), ret)]
     pub async fn add_tx(
         &self,
         tx: RpcTransaction,
         p2p_message_metadata: Option<BroadcastedMessageMetadata>,
     ) -> GatewayResult<TransactionHash> {
-        info!("Processing tx");
-        let blocking_task = ProcessTxBlockingTask::new(self, tx);
-        // Run the blocking task in the current span.
-        let curr_span = Span::current();
-        let add_tx_args =
-            tokio::task::spawn_blocking(move || curr_span.in_scope(|| blocking_task.process_tx()))
-                .await
-                .map_err(|join_err| {
-                    error!("Failed to process tx: {}", join_err);
-                    GatewaySpecError::UnexpectedError { data: "Internal server error".to_owned() }
-                })??;
-
+        let add_tx_args = self.business_logic.add_tx(tx).await?;
         let tx_hash = add_tx_args.tx.tx_hash();
-
         let add_tx_args = AddTransactionArgsWrapper { args: add_tx_args, p2p_message_metadata };
         self.mempool_client.add_tx(add_tx_args).await.map_err(|e| {
             error!("Failed to send tx to mempool: {}", e);
@@ -100,7 +114,7 @@ struct ProcessTxBlockingTask {
 }
 
 impl ProcessTxBlockingTask {
-    pub fn new(gateway: &Gateway, tx: RpcTransaction) -> Self {
+    pub fn new(gateway: &GatewayBusinessLogic, tx: RpcTransaction) -> Self {
         Self {
             stateless_tx_validator: gateway.stateless_tx_validator.clone(),
             stateful_tx_validator: gateway.stateful_tx_validator.clone(),
@@ -123,7 +137,7 @@ impl ProcessTxBlockingTask {
             &self.chain_info.chain_id,
         )?;
 
-        // Perfom post compilation validations.
+        // Perform post compilation validations.
         if let AccountTransaction::Declare(executable_declare_tx) = &executable_tx {
             if !executable_declare_tx.validate_compiled_class_hash() {
                 return Err(GatewaySpecError::CompiledClassHashMismatch);
diff --git a/crates/starknet_gateway/src/stateful_transaction_validator_test.rs b/crates/starknet_gateway/src/stateful_transaction_validator_test.rs
index de734b71cde..366ccd4431c 100644
--- a/crates/starknet_gateway/src/stateful_transaction_validator_test.rs
+++ b/crates/starknet_gateway/src/stateful_transaction_validator_test.rs
@@ -88,7 +88,7 @@ fn test_instantiate_validator(stateful_validator: StatefulTransactionValidator)
 
     let mut mock_state_reader_factory = MockStateReaderFactory::new();
 
-    // Make sure stateful_validator uses the latest block in the initiall call.
+    // Make sure stateful_validator uses the latest block in the initial call.
     let latest_state_reader = state_reader_factory.get_state_reader_from_latest_block();
     mock_state_reader_factory
         .expect_get_state_reader_from_latest_block()