From d74730af3b8fbbf67d8db25ae43916e67a4b89b6 Mon Sep 17 00:00:00 2001
From: Arsenii Kulikov <klkvrr@gmail.com>
Date: Mon, 28 Oct 2024 16:31:08 +0400
Subject: [PATCH] feat: add a wrapper for `BestTransactions` prioritizing given
 senders (#12123)

---
 crates/transaction-pool/src/pool/best.rs | 86 +++++++++++++++++++++++-
 crates/transaction-pool/src/pool/mod.rs  |  2 +-
 crates/transaction-pool/src/traits.rs    |  4 +-
 3 files changed, 88 insertions(+), 4 deletions(-)

diff --git a/crates/transaction-pool/src/pool/best.rs b/crates/transaction-pool/src/pool/best.rs
index 268e3e262c..763572e7e8 100644
--- a/crates/transaction-pool/src/pool/best.rs
+++ b/crates/transaction-pool/src/pool/best.rs
@@ -2,10 +2,10 @@ use crate::{
     identifier::TransactionId, pool::pending::PendingTransaction, PoolTransaction,
     TransactionOrdering, ValidPoolTransaction,
 };
-use alloy_primitives::B256 as TxHash;
+use alloy_primitives::{Address, B256 as TxHash};
 use core::fmt;
 use std::{
-    collections::{BTreeMap, BTreeSet, HashSet},
+    collections::{BTreeMap, BTreeSet, HashSet, VecDeque},
     sync::Arc,
 };
 
@@ -259,6 +259,88 @@ impl<I: fmt::Debug, P> fmt::Debug for BestTransactionFilter<I, P> {
     }
 }
 
+/// Wrapper over [`crate::traits::BestTransactions`] that prioritizes transactions of certain
+/// senders capping total gas used by such transactions.
+#[derive(Debug)]
+pub struct BestTransactionsWithPrioritizedSenders<I: Iterator> {
+    /// Inner iterator
+    inner: I,
+    /// A set of senders which transactions should be prioritized
+    prioritized_senders: HashSet<Address>,
+    /// Maximum total gas limit of prioritized transactions
+    max_prioritized_gas: u64,
+    /// Buffer with transactions that are not being prioritized. Those will be the first to be
+    /// included after the prioritized transactions
+    buffer: VecDeque<I::Item>,
+    /// Tracker of total gas limit of prioritized transactions. Once it reaches
+    /// `max_prioritized_gas` no more transactions will be prioritized
+    prioritized_gas: u64,
+}
+
+impl<I: Iterator> BestTransactionsWithPrioritizedSenders<I> {
+    /// Constructs a new [`BestTransactionsWithPrioritizedSenders`].
+    pub fn new(prioritized_senders: HashSet<Address>, max_prioritized_gas: u64, inner: I) -> Self {
+        Self {
+            inner,
+            prioritized_senders,
+            max_prioritized_gas,
+            buffer: Default::default(),
+            prioritized_gas: Default::default(),
+        }
+    }
+}
+
+impl<I, T> Iterator for BestTransactionsWithPrioritizedSenders<I>
+where
+    I: crate::traits::BestTransactions<Item = Arc<ValidPoolTransaction<T>>>,
+    T: PoolTransaction,
+{
+    type Item = <I as Iterator>::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // If we have space, try prioritizing transactions
+        if self.prioritized_gas < self.max_prioritized_gas {
+            for item in &mut self.inner {
+                if self.prioritized_senders.contains(&item.transaction.sender()) &&
+                    self.prioritized_gas + item.transaction.gas_limit() <=
+                        self.max_prioritized_gas
+                {
+                    self.prioritized_gas += item.transaction.gas_limit();
+                    return Some(item)
+                }
+                self.buffer.push_back(item);
+            }
+        }
+
+        if let Some(item) = self.buffer.pop_front() {
+            Some(item)
+        } else {
+            self.inner.next()
+        }
+    }
+}
+
+impl<I, T> crate::traits::BestTransactions for BestTransactionsWithPrioritizedSenders<I>
+where
+    I: crate::traits::BestTransactions<Item = Arc<ValidPoolTransaction<T>>>,
+    T: PoolTransaction,
+{
+    fn mark_invalid(&mut self, tx: &Self::Item) {
+        self.inner.mark_invalid(tx)
+    }
+
+    fn no_updates(&mut self) {
+        self.inner.no_updates()
+    }
+
+    fn set_skip_blobs(&mut self, skip_blobs: bool) {
+        if skip_blobs {
+            self.buffer.retain(|tx| !tx.transaction.is_eip4844())
+        }
+        self.inner.set_skip_blobs(skip_blobs)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/crates/transaction-pool/src/pool/mod.rs b/crates/transaction-pool/src/pool/mod.rs
index a408c76841..69f17504f7 100644
--- a/crates/transaction-pool/src/pool/mod.rs
+++ b/crates/transaction-pool/src/pool/mod.rs
@@ -106,7 +106,7 @@ use crate::{
     traits::{GetPooledTransactionLimit, NewBlobSidecar, TransactionListenerKind},
     validate::ValidTransaction,
 };
-pub use best::BestTransactionFilter;
+pub use best::{BestTransactionFilter, BestTransactionsWithPrioritizedSenders};
 pub use blob::{blob_tx_priority, fee_delta};
 pub use events::{FullTransactionEvent, TransactionEvent};
 pub use listener::{AllTransactionsEvents, TransactionEvents};
diff --git a/crates/transaction-pool/src/traits.rs b/crates/transaction-pool/src/traits.rs
index ff6e8855ef..fbbddb98f4 100644
--- a/crates/transaction-pool/src/traits.rs
+++ b/crates/transaction-pool/src/traits.rs
@@ -776,7 +776,9 @@ pub trait BestTransactions: Iterator + Send {
     /// If called then the iterator will no longer yield blob transactions.
     ///
     /// Note: this will also exclude any transactions that depend on blob transactions.
-    fn skip_blobs(&mut self);
+    fn skip_blobs(&mut self) {
+        self.set_skip_blobs(true);
+    }
 
     /// Controls whether the iterator skips blob transactions or not.
     ///