feat: expose pool transaction in PayloadTransactions (#14249)

Co-authored-by: Hamdi Allam <hamdi.allam97@gmail.com>

crates/payload/util/Cargo.toml

@@ -13,7 +13,7 @@ workspace = true
 
 [dependencies]
 # reth
-reth-primitives.workspace = true
+reth-transaction-pool.workspace = true
 
 # alloy
 alloy-primitives.workspace = true

crates/payload/util/src/lib.rs

@@ -11,5 +11,5 @@
 mod traits;
 mod transaction;
 
-pub use traits::{NoopPayloadTransactions, PayloadTransactions};
+pub use traits::{BestPayloadTransactions, NoopPayloadTransactions, PayloadTransactions};
 pub use transaction::{PayloadTransactionsChain, PayloadTransactionsFixed};

crates/payload/util/src/traits.rs

@@ -1,5 +1,7 @@
-use alloy_primitives::Address;
-use reth_primitives::Recovered;
+use std::sync::Arc;
+
+use alloy_primitives::{map::HashSet, Address};
+use reth_transaction_pool::{PoolTransaction, ValidPoolTransaction};
 
 /// Iterator that returns transactions for the block building process in the order they should be
 /// included in the block.
@@ -15,7 +17,7 @@ pub trait PayloadTransactions {
         &mut self,
         // In the future, `ctx` can include access to state for block building purposes.
         ctx: (),
-    ) -> Option<Recovered<Self::Transaction>>;
+    ) -> Option<Self::Transaction>;
 
     /// Exclude descendants of the transaction with given sender and nonce from the iterator,
     /// because this transaction won't be included in the block.
@@ -35,9 +37,149 @@ impl<T> Default for NoopPayloadTransactions<T> {
 impl<T> PayloadTransactions for NoopPayloadTransactions<T> {
     type Transaction = T;
 
-    fn next(&mut self, _ctx: ()) -> Option<Recovered<Self::Transaction>> {
+    fn next(&mut self, _ctx: ()) -> Option<Self::Transaction> {
         None
     }
 
     fn mark_invalid(&mut self, _sender: Address, _nonce: u64) {}
 }
+
+/// Wrapper struct that allows to convert `BestTransactions` (used in tx pool) to
+/// `PayloadTransactions` (used in block composition).
+#[derive(Debug)]
+pub struct BestPayloadTransactions<T, I>
+where
+    T: PoolTransaction,
+    I: Iterator<Item = Arc<ValidPoolTransaction<T>>>,
+{
+    invalid: HashSet<Address>,
+    best: I,
+}
+
+impl<T, I> BestPayloadTransactions<T, I>
+where
+    T: PoolTransaction,
+    I: Iterator<Item = Arc<ValidPoolTransaction<T>>>,
+{
+    /// Create a new `BestPayloadTransactions` with the given iterator.
+    pub fn new(best: I) -> Self {
+        Self { invalid: Default::default(), best }
+    }
+}
+
+impl<T, I> PayloadTransactions for BestPayloadTransactions<T, I>
+where
+    T: PoolTransaction,
+    I: Iterator<Item = Arc<ValidPoolTransaction<T>>>,
+{
+    type Transaction = T;
+
+    fn next(&mut self, _ctx: ()) -> Option<Self::Transaction> {
+        loop {
+            let tx = self.best.next()?;
+            if self.invalid.contains(&tx.sender()) {
+                continue
+            }
+            return Some(tx.transaction.clone())
+        }
+    }
+
+    fn mark_invalid(&mut self, sender: Address, _nonce: u64) {
+        self.invalid.insert(sender);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use crate::{
+        BestPayloadTransactions, PayloadTransactions, PayloadTransactionsChain,
+        PayloadTransactionsFixed,
+    };
+    use alloy_primitives::{map::HashSet, Address};
+    use reth_transaction_pool::{
+        pool::{BestTransactionsWithPrioritizedSenders, PendingPool},
+        test_utils::{MockOrdering, MockTransaction, MockTransactionFactory},
+        PoolTransaction,
+    };
+
+    #[test]
+    fn test_best_transactions_chained_iterators() {
+        let mut priority_pool = PendingPool::new(MockOrdering::default());
+        let mut pool = PendingPool::new(MockOrdering::default());
+        let mut f = MockTransactionFactory::default();
+
+        // Block composition
+        // ===
+        // (1) up to 100 gas: custom top-of-block transaction
+        // (2) up to 100 gas: transactions from the priority pool
+        // (3) up to 200 gas: only transactions from address A
+        // (4) up to 200 gas: only transactions from address B
+        // (5) until block gas limit: all transactions from the main pool
+
+        // Notes:
+        // - If prioritized addresses overlap, a single transaction will be prioritized twice and
+        //   therefore use the per-segment gas limit twice.
+        // - Priority pool and main pool must synchronize between each other to make sure there are
+        //   no conflicts for the same nonce. For example, in this scenario, pools can't reject
+        //   transactions with seemingly incorrect nonces, because previous transactions might be in
+        //   the other pool.
+
+        let address_top_of_block = Address::random();
+        let address_in_priority_pool = Address::random();
+        let address_a = Address::random();
+        let address_b = Address::random();
+        let address_regular = Address::random();
+
+        // Add transactions to the main pool
+        {
+            let prioritized_tx_a =
+                MockTransaction::eip1559().with_gas_price(5).with_sender(address_a);
+            // without our custom logic, B would be prioritized over A due to gas price:
+            let prioritized_tx_b =
+                MockTransaction::eip1559().with_gas_price(10).with_sender(address_b);
+            let regular_tx =
+                MockTransaction::eip1559().with_gas_price(15).with_sender(address_regular);
+            pool.add_transaction(Arc::new(f.validated(prioritized_tx_a)), 0);
+            pool.add_transaction(Arc::new(f.validated(prioritized_tx_b)), 0);
+            pool.add_transaction(Arc::new(f.validated(regular_tx)), 0);
+        }
+
+        // Add transactions to the priority pool
+        {
+            let prioritized_tx =
+                MockTransaction::eip1559().with_gas_price(0).with_sender(address_in_priority_pool);
+            let valid_prioritized_tx = f.validated(prioritized_tx);
+            priority_pool.add_transaction(Arc::new(valid_prioritized_tx), 0);
+        }
+
+        let mut block = PayloadTransactionsChain::new(
+            PayloadTransactionsFixed::single(
+                MockTransaction::eip1559().with_sender(address_top_of_block),
+            ),
+            Some(100),
+            PayloadTransactionsChain::new(
+                BestPayloadTransactions::new(priority_pool.best()),
+                Some(100),
+                BestPayloadTransactions::new(BestTransactionsWithPrioritizedSenders::new(
+                    HashSet::from([address_a]),
+                    200,
+                    BestTransactionsWithPrioritizedSenders::new(
+                        HashSet::from([address_b]),
+                        200,
+                        pool.best(),
+                    ),
+                )),
+                None,
+            ),
+            None,
+        );
+
+        assert_eq!(block.next(()).unwrap().sender(), address_top_of_block);
+        assert_eq!(block.next(()).unwrap().sender(), address_in_priority_pool);
+        assert_eq!(block.next(()).unwrap().sender(), address_a);
+        assert_eq!(block.next(()).unwrap().sender(), address_b);
+        assert_eq!(block.next(()).unwrap().sender(), address_regular);
+    }
+}

crates/payload/util/src/transaction.rs

@@ -1,7 +1,7 @@
 use crate::PayloadTransactions;
 use alloy_consensus::Transaction;
 use alloy_primitives::Address;
-use reth_primitives::Recovered;
+use reth_transaction_pool::PoolTransaction;
 
 /// An implementation of [`crate::traits::PayloadTransactions`] that yields
 /// a pre-defined set of transactions.
@@ -26,10 +26,10 @@ impl<T> PayloadTransactionsFixed<T> {
     }
 }
 
-impl<T: Clone> PayloadTransactions for PayloadTransactionsFixed<Recovered<T>> {
+impl<T: Clone> PayloadTransactions for PayloadTransactionsFixed<T> {
     type Transaction = T;
 
-    fn next(&mut self, _ctx: ()) -> Option<Recovered<T>> {
+    fn next(&mut self, _ctx: ()) -> Option<T> {
         (self.index < self.transactions.len()).then(|| {
             let tx = self.transactions[self.index].clone();
             self.index += 1;
@@ -91,20 +91,20 @@ impl<B: PayloadTransactions, A: PayloadTransactions> PayloadTransactionsChain<B,
 
 impl<A, B> PayloadTransactions for PayloadTransactionsChain<A, B>
 where
-    A: PayloadTransactions<Transaction: Transaction>,
+    A: PayloadTransactions<Transaction: PoolTransaction>,
     B: PayloadTransactions<Transaction = A::Transaction>,
 {
     type Transaction = A::Transaction;
 
-    fn next(&mut self, ctx: ()) -> Option<Recovered<Self::Transaction>> {
+    fn next(&mut self, ctx: ()) -> Option<Self::Transaction> {
         while let Some(tx) = self.before.next(ctx) {
             if let Some(before_max_gas) = self.before_max_gas {
-                if self.before_gas + tx.tx().gas_limit() <= before_max_gas {
-                    self.before_gas += tx.tx().gas_limit();
+                if self.before_gas + tx.gas_limit() <= before_max_gas {
+                    self.before_gas += tx.gas_limit();
                     return Some(tx);
                 }
-                self.before.mark_invalid(tx.signer(), tx.tx().nonce());
-                self.after.mark_invalid(tx.signer(), tx.tx().nonce());
+                self.before.mark_invalid(tx.sender(), tx.nonce());
+                self.after.mark_invalid(tx.sender(), tx.nonce());
             } else {
                 return Some(tx);
             }
@@ -112,11 +112,11 @@ where
 
         while let Some(tx) = self.after.next(ctx) {
             if let Some(after_max_gas) = self.after_max_gas {
-                if self.after_gas + tx.tx().gas_limit() <= after_max_gas {
-                    self.after_gas += tx.tx().gas_limit();
+                if self.after_gas + tx.gas_limit() <= after_max_gas {
+                    self.after_gas += tx.gas_limit();
                     return Some(tx);
                 }
-                self.after.mark_invalid(tx.signer(), tx.tx().nonce());
+                self.after.mark_invalid(tx.sender(), tx.nonce());
             } else {
                 return Some(tx);
             }