txpool reload queued senders on usdc storage changes

Author: drappi-ai

crates/seismic/txpool/src/maintain.rs

@@ -1,9 +1,11 @@
 //! Seismic-specific pool maintenance hook that augments native balances with
 //! USDC predeploy balances.
 
-use reth_execution_types::ChangedAccount;
+use alloy_primitives::{Address, B256};
+use reth_execution_types::{ChangedAccount, ExecutionOutcome};
 use reth_provider::StateProviderFactory;
 use reth_transaction_pool::maintain::ChangedAccountsHook;
+use std::collections::HashSet;
 use tracing::{debug, warn};
 
 /// A [`ChangedAccountsHook`] that reads each sender's USDC predeploy balance
@@ -50,4 +52,122 @@ where
             }
         }
     }
+
+    fn extend_reload_queued_senders<R>(
+        &self,
+        queued_senders: &HashSet<Address>,
+        old: Option<&ExecutionOutcome<R>>,
+        new: &ExecutionOutcome<R>,
+        dirty_addresses: &mut HashSet<Address>,
+    ) {
+        dirty_addresses.extend(queued_senders_with_changed_usdc_slots(queued_senders, old, new));
+    }
+}
+
+fn queued_senders_with_changed_usdc_slots<'a, R>(
+    queued_senders: &'a HashSet<Address>,
+    old: Option<&ExecutionOutcome<R>>,
+    new: &ExecutionOutcome<R>,
+) -> impl Iterator<Item = Address> + 'a {
+    let changed_slots = changed_usdc_storage_slots(new)
+        .into_iter()
+        .chain(old.into_iter().flat_map(changed_usdc_storage_slots))
+        .collect::<HashSet<_>>();
+
+    queued_senders.iter().copied().filter(move |address| {
+        changed_slots.contains(&crate::usdc::usdc_balance_storage_key(address))
+    })
+}
+
+fn changed_usdc_storage_slots<R>(state: &ExecutionOutcome<R>) -> impl Iterator<Item = B256> + '_ {
+    state
+        .bundle_accounts_iter()
+        .filter_map(|(address, account)| (address == crate::usdc::USDC_CONTRACT).then_some(account))
+        .flat_map(|account| {
+            account.storage.iter().filter_map(|(slot, value)| {
+                value.is_changed().then(|| B256::from(slot.to_be_bytes::<32>()))
+            })
+        })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use alloy_primitives::{map::HashMap, FlaggedStorage, U256};
+    use reth_execution_types::{BundleStateInit, RevertsInit};
+
+    #[test]
+    fn reloads_only_queued_senders_with_changed_usdc_slots() {
+        let affected = alloy_primitives::address!("000000000000000000000000000000000000000a");
+        let unaffected = alloy_primitives::address!("000000000000000000000000000000000000000b");
+        let queued_senders = HashSet::from([affected, unaffected]);
+        let state = ExecutionOutcome::<()>::new_init(
+            {
+                let mut init = BundleStateInit::default();
+                init.insert(
+                    crate::usdc::USDC_CONTRACT,
+                    (
+                        None,
+                        None,
+                        HashMap::from_iter([(
+                            crate::usdc::usdc_balance_storage_key(&affected),
+                            (FlaggedStorage::ZERO, FlaggedStorage::from(U256::from(1))),
+                        )]),
+                    ),
+                );
+                init
+            },
+            RevertsInit::default(),
+            [],
+            vec![],
+            0,
+            vec![],
+        );
+
+        let dirty = queued_senders_with_changed_usdc_slots(&queued_senders, None, &state)
+            .collect::<HashSet<_>>();
+
+        assert_eq!(dirty, HashSet::from([affected]));
+    }
+
+    #[test]
+    fn includes_changed_slots_from_old_and_new_state() {
+        let queued_sender = alloy_primitives::address!("000000000000000000000000000000000000000a");
+        let queued_senders = HashSet::from([queued_sender]);
+        let old = ExecutionOutcome::<()>::new_init(
+            {
+                let mut init = BundleStateInit::default();
+                init.insert(
+                    crate::usdc::USDC_CONTRACT,
+                    (
+                        None,
+                        None,
+                        HashMap::from_iter([(
+                            crate::usdc::usdc_balance_storage_key(&queued_sender),
+                            (FlaggedStorage::ZERO, FlaggedStorage::from(U256::from(1))),
+                        )]),
+                    ),
+                );
+                init
+            },
+            RevertsInit::default(),
+            [],
+            vec![],
+            0,
+            vec![],
+        );
+        let new = ExecutionOutcome::<()>::new_init(
+            BundleStateInit::default(),
+            RevertsInit::default(),
+            [],
+            vec![],
+            0,
+            vec![],
+        );
+
+        let dirty = queued_senders_with_changed_usdc_slots(&queued_senders, Some(&old), &new)
+            .collect::<HashSet<_>>();
+
+        assert_eq!(dirty, HashSet::from([queued_sender]));
+    }
 }

crates/seismic/txpool/src/usdc.rs

@@ -22,7 +22,7 @@ const BALANCES_MAPPING_SLOT: U256 = U256::from_limbs([3, 0, 0, 0]);
 /// For a Solidity `mapping(address => uint256)` at slot `s`, the value for key
 /// `k` is stored at `keccak256(abi.encode(k, s))` — i.e. `k` left-padded to 32
 /// bytes concatenated with `s` as a 32-byte big-endian integer.
-fn usdc_balance_storage_key(address: &Address) -> B256 {
+pub(crate) fn usdc_balance_storage_key(address: &Address) -> B256 {
     let mut buf = [0u8; 64];
     // address is 20 bytes, right-aligned in the first 32-byte word
     buf[12..32].copy_from_slice(address.as_slice());

crates/transaction-pool/src/maintain.rs

@@ -17,7 +17,7 @@ use futures_util::{
 };
 use reth_chain_state::CanonStateNotification;
 use reth_chainspec::{ChainSpecProvider, EthChainSpec};
-use reth_execution_types::ChangedAccount;
+use reth_execution_types::{ChangedAccount, ExecutionOutcome};
 use reth_fs_util::FsPathError;
 use reth_primitives_traits::{
     transaction::signed::SignedTransaction, NodePrimitives, SealedHeader,
@@ -100,6 +100,17 @@ pub trait ChangedAccountsHook: Send + Sync + 'static {
     /// additional state (e.g. ERC-20 storage) and adjust the `balance` field of
     /// each [`ChangedAccount`].
     fn transform(&self, accounts: &mut Vec<ChangedAccount>);
+
+    /// Extends the dirty set with queued senders whose executability may have
+    /// changed due to off-account state updates in the canonical update.
+    fn extend_reload_queued_senders<R>(
+        &self,
+        _queued_senders: &HashSet<Address>,
+        _old: Option<&ExecutionOutcome<R>>,
+        _new: &ExecutionOutcome<R>,
+        _dirty_addresses: &mut HashSet<Address>,
+    ) {
+    }
 }
 
 /// No-op implementation for chains that don't need balance augmentation.
@@ -383,8 +394,8 @@ pub async fn maintain_transaction_pool_with_hook<N, Client, P, St, Tasks, H>(
                 let old_first = old_blocks.first();
 
                 // check if the reorg is not canonical with the pool's block
-                if !(old_first.parent_hash() == pool_info.last_seen_block_hash ||
-                    new_first.parent_hash() == pool_info.last_seen_block_hash)
+                if !(old_first.parent_hash() == pool_info.last_seen_block_hash
+                    || new_first.parent_hash() == pool_info.last_seen_block_hash)
                 {
                     // the new block points to a higher block than the oldest block in the old chain
                     maintained_state = MaintainedPoolState::Drifted;
@@ -436,7 +447,6 @@ pub async fn maintain_transaction_pool_with_hook<N, Client, P, St, Tasks, H>(
                 // we can use extend here because they are unique
                 changed_accounts.extend(new_changed_accounts.into_iter().map(|entry| entry.0));
                 hook.transform(&mut changed_accounts);
-
                 // all transactions mined in the new chain
                 let new_mined_transactions: HashSet<_> = new_blocks.transaction_hashes().collect();
 
@@ -488,6 +498,14 @@ pub async fn maintain_transaction_pool_with_hook<N, Client, P, St, Tasks, H>(
                 metrics.inc_reinserted_transactions(pruned_old_transactions.len());
                 let _ = pool.add_external_transactions(pruned_old_transactions).await;
 
+                let queued_senders = queued_senders(&pool);
+                hook.extend_reload_queued_senders(
+                    &queued_senders,
+                    Some(old_state),
+                    new_state,
+                    &mut dirty_addresses,
+                );
+
                 // keep track of new mined blob transactions
                 blob_store_tracker.add_new_chain_blocks(&new_blocks);
             }
@@ -531,7 +549,7 @@ pub async fn maintain_transaction_pool_with_hook<N, Client, P, St, Tasks, H>(
                     // keep track of mined blob transactions
                     blob_store_tracker.add_new_chain_blocks(&blocks);
 
-                    continue
+                    continue;
                 }
 
                 let mut changed_accounts = Vec::with_capacity(state.state().len());
@@ -541,7 +559,6 @@ pub async fn maintain_transaction_pool_with_hook<N, Client, P, St, Tasks, H>(
                     changed_accounts.push(acc);
                 }
                 hook.transform(&mut changed_accounts);
-
                 let mined_transactions = blocks.transaction_hashes().collect();
 
                 // check if the range of the commit is canonical with the pool's block
@@ -563,6 +580,14 @@ pub async fn maintain_transaction_pool_with_hook<N, Client, P, St, Tasks, H>(
                 };
                 pool.on_canonical_state_change(update);
 
+                let queued_senders = queued_senders(&pool);
+                hook.extend_reload_queued_senders(
+                    &queued_senders,
+                    None,
+                    state,
+                    &mut dirty_addresses,
+                );
+
                 // keep track of mined blob transactions
                 blob_store_tracker.add_new_chain_blocks(&blocks);
             }
@@ -671,6 +696,10 @@ where
     Ok(res)
 }
 
+fn queued_senders<P: TransactionPool>(pool: &P) -> HashSet<Address> {
+    pool.queued_transactions().into_iter().map(|tx| tx.sender()).collect()
+}
+
 /// Loads transactions from a file, decodes them from the JSON or RLP format, and
 /// inserts them into the transaction pool on node boot up.
 /// The file is removed after the transactions have been successfully processed.
@@ -682,14 +711,14 @@ where
     P: TransactionPool<Transaction: PoolTransaction<Consensus: SignedTransaction>>,
 {
     if !file_path.exists() {
-        return Ok(())
+        return Ok(());
     }
 
     debug!(target: "txpool", txs_file =?file_path, "Check local persistent storage for saved transactions");
     let data = reth_fs_util::read(file_path)?;
 
     if data.is_empty() {
-        return Ok(())
+        return Ok(());
     }
 
     let pool_transactions: Vec<(TransactionOrigin, <P as TransactionPool>::Transaction)> =
@@ -740,7 +769,7 @@ where
     let local_transactions = pool.get_local_transactions();
     if local_transactions.is_empty() {
         trace!(target: "txpool", "no local transactions to save");
-        return
+        return;
     }
 
     let local_transactions = local_transactions
@@ -757,7 +786,7 @@ where
         Ok(data) => data,
         Err(err) => {
             warn!(target: "txpool", %err, txs_file=?file_path, "failed to serialize local transactions to json");
-            return
+            return;
         }
     };
 
@@ -812,7 +841,7 @@ pub async fn backup_local_transactions_task<P>(
 {
     let Some(transactions_path) = config.transactions_path else {
         // nothing to do
-        return
+        return;
     };
 
     if let Err(err) = load_and_reinsert_transactions(pool.clone(), &transactions_path).await {