fix(concurrency): add the sequencer key when checking if the tx is to…

…o big in concurrency mode

crates/blockifier/src/concurrency/worker_logic.rs

@@ -1,6 +1,6 @@
 use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
-use std::sync::{Arc, Mutex};
+use std::sync::Mutex;
 
 use num_traits::ToPrimitive;
 use starknet_api::core::{ClassHash, ContractAddress};
@@ -223,12 +223,14 @@ impl<'a, S: StateReader> WorkerExecutor<'a, S> {
         let mut execution_output = lock_mutex_in_array(&self.execution_outputs, tx_index);
         let writes = &execution_output.as_ref().expect(EXECUTION_OUTPUTS_UNWRAP_ERROR).writes;
         let reads = &execution_output.as_ref().expect(EXECUTION_OUTPUTS_UNWRAP_ERROR).reads;
-        let tx_state_changes_keys = StateChanges::from(writes.diff(reads)).into_keys();
+        let mut tx_state_changes_keys = StateChanges::from(writes.diff(reads)).into_keys();
+        let tx_context = self.block_context.to_tx_context(tx);
         let tx_result =
             &mut execution_output.as_mut().expect(EXECUTION_OUTPUTS_UNWRAP_ERROR).result;
 
-        let tx_context = Arc::new(self.block_context.to_tx_context(tx));
         if let Ok(tx_execution_info) = tx_result.as_mut() {
+            // Add the deleted sequencer balance key to the storage keys.
+            tx_state_changes_keys.update_sequencer_key_in_storage(&tx_context, tx_execution_info);
             // Ask the bouncer if there is room for the transaction in the block.
             let bouncer_result = self.bouncer.lock().expect("Bouncer lock failed.").try_update(
                 &tx_versioned_state,
@@ -246,9 +248,7 @@ impl<'a, S: StateReader> WorkerExecutor<'a, S> {
                 }
             }
             // Update the sequencer balance (in state + call info).
-            if tx_context.tx_info.sender_address()
-                == self.block_context.block_info.sequencer_address
-            {
+            if tx_context.is_sequencer_the_sender() {
                 // When the sequencer is the sender, we use the sequential (full) fee transfer.
                 return true;
             }

crates/blockifier/src/context.rs

@@ -18,6 +18,9 @@ impl TransactionContext {
     pub fn fee_token_address(&self) -> ContractAddress {
         self.block_context.chain_info.fee_token_address(&self.tx_info.fee_type())
     }
+    pub fn is_sequencer_the_sender(&self) -> bool {
+        self.tx_info.sender_address() == self.block_context.block_info.sequencer_address
+    }
 }
 
 #[derive(Clone, Debug)]

crates/blockifier/src/state/cached_state.rs

@@ -8,9 +8,11 @@ use starknet_api::hash::StarkFelt;
 use starknet_api::state::StorageKey;
 
 use crate::abi::abi_utils::get_fee_token_var_address;
+use crate::context::TransactionContext;
 use crate::execution::contract_class::ContractClass;
 use crate::state::errors::StateError;
 use crate::state::state_api::{State, StateReader, StateResult, UpdatableState};
+use crate::transaction::objects::TransactionExecutionInfo;
 use crate::utils::{strict_subtract_mappings, subtract_mappings};
 
 #[cfg(test)]
@@ -617,6 +619,23 @@ impl StateChangesKeys {
         self.modified_contracts.extend(&other.modified_contracts);
     }
 
+    pub fn update_sequencer_key_in_storage(
+        &mut self,
+        tx_context: &TransactionContext,
+        tx_result: &TransactionExecutionInfo,
+    ) {
+        let actual_fee = tx_result.transaction_receipt.fee.0;
+        let sequencer_address = tx_context.block_context.block_info.sequencer_address;
+        if tx_context.block_context.concurrency_mode
+            && !tx_context.is_sequencer_the_sender()
+            && actual_fee > 0
+        {
+            // Add the deleted sequencer balance key to the storage keys.
+            let sequencer_balance_low = get_fee_token_var_address(sequencer_address);
+            self.storage_keys.insert((tx_context.fee_token_address(), sequencer_balance_low));
+        }
+    }
+
     pub fn count(&self) -> StateChangesCount {
         // nonce_keys effect is captured by modified_contracts; it is not used but kept for
         // completeness of this struct.

crates/blockifier/src/transaction/account_transaction.rs

@@ -318,14 +318,12 @@ impl AccountTransaction {
         // TODO(Amos, 8/04/2024): Add test for this assert.
         Self::assert_actual_fee_in_bounds(&tx_context, actual_fee)?;
 
-        let fee_transfer_call_info = if tx_context.block_context.concurrency_mode
-            && tx_context.block_context.block_info.sequencer_address
-                != tx_context.tx_info.sender_address()
-        {
-            Self::concurrency_execute_fee_transfer(state, tx_context, actual_fee)?
-        } else {
-            Self::execute_fee_transfer(state, tx_context, actual_fee)?
-        };
+        let fee_transfer_call_info =
+            if tx_context.block_context.concurrency_mode && !tx_context.is_sequencer_the_sender() {
+                Self::concurrency_execute_fee_transfer(state, tx_context, actual_fee)?
+            } else {
+                Self::execute_fee_transfer(state, tx_context, actual_fee)?
+            };
 
         Ok(Some(fee_transfer_call_info))
     }

crates/blockifier/src/transaction/transaction_execution.rs

@@ -175,7 +175,11 @@ impl<U: UpdatableState> ExecutableTransaction<U> for Transaction {
         // Check if the transaction is too large to fit any block.
         // TODO(Yoni, 1/8/2024): consider caching these two.
         let tx_execution_summary = tx_execution_info.summarize();
-        let tx_state_changes_keys = state.get_actual_state_changes()?.into_keys();
+        let mut tx_state_changes_keys = state.get_actual_state_changes()?.into_keys();
+        tx_state_changes_keys.update_sequencer_key_in_storage(
+            &block_context.to_tx_context(self),
+            &tx_execution_info,
+        );
         verify_tx_weights_in_bounds(
             state,
             &tx_execution_summary,