diff --git a/crates/mempool/src/mempool.rs b/crates/mempool/src/mempool.rs index d6ee6382..ce529d7e 100644 --- a/crates/mempool/src/mempool.rs +++ b/crates/mempool/src/mempool.rs @@ -54,6 +54,7 @@ impl Mempool { let tx = self.tx_pool.remove(tx_hash)?; eligible_txs.push(tx); } + self.enqueue_next_eligible_txs(&eligible_txs)?; Ok(eligible_txs) } @@ -111,6 +112,21 @@ impl Mempool { Ok(()) } + fn enqueue_next_eligible_txs(&mut self, txs: &[ThinTransaction]) -> MempoolResult<()> { + for tx in txs { + let current_account_state = Account { + sender_address: tx.sender_address, + state: AccountState { nonce: tx.nonce }, + }; + if let Some(next_tx_reference) = + self.tx_pool.get_next_eligible_tx(current_account_state)? + { + self.tx_queue.insert(*next_tx_reference); + } + } + Ok(()) + } + #[cfg(test)] pub(crate) fn _tx_pool(&self) -> &TransactionPool { &self.tx_pool diff --git a/crates/mempool/src/mempool_test.rs b/crates/mempool/src/mempool_test.rs index 8f015548..17cc3039 100644 --- a/crates/mempool/src/mempool_test.rs +++ b/crates/mempool/src/mempool_test.rs @@ -175,8 +175,6 @@ fn test_get_txs(#[case] requested_txs: usize) { } #[rstest] -// TODO(Ayelet): remove ignore once replenishing is merged. -#[ignore] fn test_get_txs_multi_nonce() { // Setup. let tx_address_0_nonce_0 = @@ -192,8 +190,11 @@ fn test_get_txs_multi_nonce() { let txs = mempool.get_txs(2).unwrap(); // Assert that the account's next tx was added the queue. - assert_eq!(txs, &[tx_address_0_nonce_0, tx_address_0_nonce_1]); - let expected_mempool_state = MempoolState::new([], []); + // TODO(Ayelet): all transactions should be returned after replenishing. + assert_eq!(txs, &[tx_address_0_nonce_0]); + let expected_queue_txs = [TransactionReference::new(&tx_address_0_nonce_1)]; + let expected_pool_txs = [tx_address_0_nonce_1]; + let expected_mempool_state = MempoolState::new(expected_pool_txs, expected_queue_txs); expected_mempool_state.assert_eq_mempool_state(&mempool); } diff --git a/crates/mempool/src/transaction_pool.rs b/crates/mempool/src/transaction_pool.rs index 3c59f07c..27610c57 100644 --- a/crates/mempool/src/transaction_pool.rs +++ b/crates/mempool/src/transaction_pool.rs @@ -3,7 +3,9 @@ use std::collections::{hash_map, BTreeMap, HashMap}; use starknet_api::core::{ContractAddress, Nonce}; use starknet_api::transaction::TransactionHash; use starknet_mempool_types::errors::MempoolError; -use starknet_mempool_types::mempool_types::{MempoolResult, ThinTransaction}; +use starknet_mempool_types::mempool_types::{ + Account, AccountState, MempoolResult, ThinTransaction, +}; use crate::mempool::TransactionReference; @@ -86,6 +88,16 @@ impl TransactionPool { self.txs_by_account.get(address, nonce) } + pub fn get_next_eligible_tx( + &self, + current_account_state: Account, + ) -> MempoolResult> { + let Account { sender_address, state: AccountState { nonce } } = current_account_state; + // TOOD(Ayelet): Change to StarknetApiError. + let next_nonce = nonce.try_increment().map_err(|_| MempoolError::FeltOutOfRange)?; + Ok(self.get_by_address_and_nonce(sender_address, next_nonce)) + } + #[cfg(test)] pub(crate) fn _tx_pool(&self) -> &HashToTransaction { &self.tx_pool