From 12c2a87bda511868d7000cf9316f92c4923e4aed Mon Sep 17 00:00:00 2001 From: sword-smith Date: Tue, 15 Oct 2024 22:41:57 +0200 Subject: [PATCH] mempool: allow capping number of txs in mempool Updating transaction proofs is a very expensive procedure. Capping the number of transactions in the mempool addresses this problem. This cap allows miners to work on the block proof without waiting too long for the updating of the transaction proofs. --- src/config_models/cli_args.rs | 10 ++++ src/lib.rs | 6 ++- src/models/state/mempool.rs | 88 ++++++++++++++++++++++++++++++++--- src/tests/shared.rs | 7 ++- 4 files changed, 102 insertions(+), 9 deletions(-) diff --git a/src/config_models/cli_args.rs b/src/config_models/cli_args.rs index 724ba25b..8bd30a25 100644 --- a/src/config_models/cli_args.rs +++ b/src/config_models/cli_args.rs @@ -69,6 +69,15 @@ pub struct Args { #[clap(long, default_value = "1G", value_name = "SIZE")] pub max_mempool_size: ByteSize, + /// Maximum number of transactions permitted in the mempool. + /// + /// If too much time is spent updating transaction proofs, this + /// value can be capped. + /// + /// E.g. --max-mempool-num-tx=4 + #[clap(long)] + pub max_mempool_num_tx: Option, + /// Port on which to listen for peer connections. #[clap(long, default_value = "9798", value_name = "PORT")] pub peer_port: u16, @@ -160,6 +169,7 @@ mod cli_args_tests { IpAddr::from(Ipv6Addr::UNSPECIFIED), default_args.listen_addr ); + assert_eq!(None, default_args.max_mempool_num_tx); } #[test] diff --git a/src/lib.rs b/src/lib.rs index 12332176..eefa9d7e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -159,7 +159,11 @@ pub async fn initialize(cli_args: cli_args::Args) -> Result<()> { archival_state, }; let blockchain_state = BlockchainState::Archival(blockchain_archival_state); - let mempool = Mempool::new(cli_args.max_mempool_size, latest_block.hash()); + let mempool = Mempool::new( + cli_args.max_mempool_size, + cli_args.max_mempool_num_tx, + latest_block.hash(), + ); let mut global_state_lock = GlobalStateLock::new( wallet_state, blockchain_state, diff --git a/src/models/state/mempool.rs b/src/models/state/mempool.rs index b7adfbf7..9aca3ed2 100644 --- a/src/models/state/mempool.rs +++ b/src/models/state/mempool.rs @@ -62,8 +62,13 @@ type LookupItem<'a> = (TransactionKernelId, &'a Transaction); #[derive(Debug, Clone, PartialEq, Eq, GetSize)] pub struct Mempool { + /// Maximum size this data structure may take up in memory. max_total_size: usize, + /// If set, represents the maximum number of transactions allowed in the + /// mempool. If None, mempool is only restricted by size. + max_length: Option, + /// Contains transactions, with a mapping from transaction ID to transaction. /// Maintain for constant lookup tx_dictionary: HashMap, @@ -80,12 +85,17 @@ pub struct Mempool { impl Mempool { /// instantiate a new, empty `Mempool` - pub fn new(max_total_size: ByteSize, tip_digest: Digest) -> Self { + pub fn new( + max_total_size: ByteSize, + max_num_transactions: Option, + tip_digest: Digest, + ) -> Self { let table = Default::default(); let queue = Default::default(); let max_total_size = max_total_size.0.try_into().unwrap(); Self { max_total_size, + max_length: max_num_transactions, tx_dictionary: table, queue, tip_digest, @@ -180,6 +190,7 @@ impl Mempool { "mempool's table and queue length must agree prior to shrink" ); self.shrink_to_max_size(); + self.shrink_to_max_length(); assert_eq!( self.tx_dictionary.len(), self.queue.len(), @@ -404,7 +415,7 @@ impl Mempool { } /// Shrink the memory pool to the value of its `max_size` field. - /// Likely computes in O(n) + /// Likely computes in O(n). fn shrink_to_max_size(&mut self) { // Repeately remove the least valuable transaction while self.get_size() > self.max_total_size && self.pop_min().is_some() { @@ -414,6 +425,18 @@ impl Mempool { self.shrink_to_fit() } + /// Shrink the memory pool to the value of its `max_length` field, + /// if that field is set. + fn shrink_to_max_length(&mut self) { + if let Some(max_length) = self.max_length { + while self.len() > max_length && self.pop_min().is_some() { + continue; + } + } + + self.shrink_to_fit() + } + /// Shrinks internal data structures as much as possible. /// Computes in O(n) (Likely) fn shrink_to_fit(&mut self) { @@ -488,7 +511,7 @@ mod tests { pub async fn insert_then_get_then_remove_then_get() { let network = Network::Main; let genesis_block = Block::genesis_block(network); - let mut mempool = Mempool::new(ByteSize::gb(1), genesis_block.hash()); + let mut mempool = Mempool::new(ByteSize::gb(1), None, genesis_block.hash()); let txs = make_plenty_mock_transaction_with_primitive_witness(2); let transaction_digests = txs.iter().map(|tx| tx.kernel.txid()).collect_vec(); @@ -527,7 +550,7 @@ mod tests { /// Create a mempool with n transactions. async fn setup_mock_mempool(transactions_count: usize, network: Network) -> Mempool { let genesis_block = Block::genesis_block(network); - let mut mempool = Mempool::new(ByteSize::gb(1), genesis_block.hash()); + let mut mempool = Mempool::new(ByteSize::gb(1), None, genesis_block.hash()); let txs = make_plenty_mock_transaction_with_primitive_witness(transactions_count); for tx in txs { mempool.insert(&tx); @@ -576,7 +599,7 @@ mod tests { async fn prune_stale_transactions() { let network = Network::Main; let genesis_block = Block::genesis_block(network); - let mut mempool = Mempool::new(ByteSize::gb(1), genesis_block.hash()); + let mut mempool = Mempool::new(ByteSize::gb(1), None, genesis_block.hash()); assert!( mempool.is_empty(), "Mempool must be empty after initialization" @@ -677,7 +700,7 @@ mod tests { bob.wallet_state.add_expected_utxos(expected_utxos).await; // Add this transaction to a mempool - let mut mempool = Mempool::new(ByteSize::gb(1), block_1.hash()); + let mut mempool = Mempool::new(ByteSize::gb(1), None, block_1.hash()); mempool.insert(&tx_by_bob); // Create another transaction that's valid to be included in block 2, but isn't actually @@ -987,6 +1010,59 @@ mod tests { } } + #[traced_test] + #[tokio::test] + async fn max_len_none() { + let network = Network::Main; + let genesis_block = Block::genesis_block(network); + let txs = make_plenty_mock_transaction_with_primitive_witness(11); + let mut mempool = Mempool::new(ByteSize::gb(1), None, genesis_block.hash()); + + for tx in txs.iter() { + mempool.insert(tx); + } + + assert_eq!( + 11, + mempool.len(), + "All transactions are inserted into mempool" + ); + } + + #[traced_test] + #[tokio::test] + async fn max_len_is_respected() { + let network = Network::Main; + let genesis_block = Block::genesis_block(network); + let txs = make_plenty_mock_transaction_with_primitive_witness(20); + + let mut expected_txs = txs.clone(); + expected_txs.sort_by_key(|x| x.fee_density()); + expected_txs.reverse(); + + for i in 0..10 { + let mut mempool = Mempool::new(ByteSize::gb(1), Some(i), genesis_block.hash()); + for tx in txs.iter() { + mempool.insert(tx); + } + + assert_eq!( + i, + mempool.len(), + "Only {i} transactions are permitted in the mempool" + ); + + let expected_txs = expected_txs.iter().take(i).cloned().collect_vec(); + + let mut mempool_iter = mempool.get_sorted_iter(); + for expected_tx in expected_txs.iter() { + let (txid, fee_density) = mempool_iter.next().unwrap(); + assert_eq!(expected_tx, mempool.get(txid).unwrap()); + assert_eq!(expected_tx.fee_density(), fee_density); + } + } + } + #[traced_test] #[tokio::test] async fn get_mempool_size() { diff --git a/src/tests/shared.rs b/src/tests/shared.rs index 9d784d66..312d29e2 100644 --- a/src/tests/shared.rs +++ b/src/tests/shared.rs @@ -11,7 +11,6 @@ use std::time::SystemTime; use anyhow::Result; use bytes::Bytes; use bytes::BytesMut; -use bytesize::ByteSize; use futures::sink; use futures::stream; use futures::task::Context; @@ -210,11 +209,15 @@ pub(crate) async fn mock_genesis_global_state( light_state, archival_state, }); - let mempool = Mempool::new(ByteSize::gb(1), genesis_block.hash()); let cli_args = cli_args::Args { network, ..Default::default() }; + let mempool = Mempool::new( + cli_args.max_mempool_size, + cli_args.max_mempool_num_tx, + genesis_block.hash(), + ); let wallet_state = mock_genesis_wallet_state(wallet, network).await;