Skip to content

Commit

Permalink
feat(mempool): get txs return requested txs after replenshing queue
Browse files Browse the repository at this point in the history
  • Loading branch information
ayeletstarkware committed Jul 18, 2024
1 parent 2d407eb commit e63abfd
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 14 deletions.
19 changes: 15 additions & 4 deletions crates/mempool/src/mempool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,11 +49,22 @@ impl Mempool {
// library.
pub fn get_txs(&mut self, n_txs: usize) -> MempoolResult<Vec<ThinTransaction>> {
let mut eligible_txs: Vec<ThinTransaction> = Vec::with_capacity(n_txs);
for tx_hash in self.tx_queue.pop_chunk(n_txs) {
let tx = self.tx_pool.remove(tx_hash)?;
eligible_txs.push(tx);
let mut remaining_txs = n_txs;

while remaining_txs > 0 && !self.tx_queue.is_empty() {
let mut popped_txs = Vec::with_capacity(remaining_txs);

for tx_hash in self.tx_queue.pop_chunk(remaining_txs) {
let tx = self.tx_pool.remove(tx_hash)?;
popped_txs.push(tx);
}

self.enqueue_next_eligible_txs(&popped_txs);

remaining_txs -= popped_txs.len();

eligible_txs.extend(popped_txs);
}
self.enqueue_next_eligible_txs(&eligible_txs);

Ok(eligible_txs)
}
Expand Down
19 changes: 9 additions & 10 deletions crates/mempool/src/mempool_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -181,20 +181,20 @@ fn test_get_txs_multi_nonce() {
add_tx_input!(tx_hash: 1, sender_address: "0x0", tx_nonce: 0_u8, account_nonce: 0_u8).tx;
let tx_address_0_nonce_1 =
add_tx_input!(tx_hash: 2, sender_address: "0x0", tx_nonce: 1_u8, account_nonce: 0_u8).tx;
let tx_address_0_nonce_2 =
add_tx_input!(tx_hash: 3, sender_address: "0x0", tx_nonce: 2_u8, account_nonce: 0_u8).tx;

let queue_txs = [TransactionReference::new(&tx_address_0_nonce_0)];
let pool_txs = [tx_address_0_nonce_0.clone(), tx_address_0_nonce_1.clone()];
let pool_txs =
[tx_address_0_nonce_0.clone(), tx_address_0_nonce_1.clone(), tx_address_0_nonce_2.clone()];
let mut mempool: Mempool = MempoolState::new(pool_txs, queue_txs).into();

// Test.
let txs = mempool.get_txs(2).unwrap();
let txs = mempool.get_txs(3).unwrap();

// Assert that the account's next tx was added the queue.
// 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);
// Assert: all transactions are returned.
assert_eq!(txs, &[tx_address_0_nonce_0, tx_address_0_nonce_1, tx_address_0_nonce_2]);
let expected_mempool_state = MempoolState::new([], []);
expected_mempool_state.assert_eq_mempool_state(&mempool);
}

Expand Down Expand Up @@ -393,7 +393,6 @@ fn test_flow_filling_holes(mut mempool: Mempool) {
add_tx(&mut mempool, &input_address_0_nonce_0);
let txs = mempool.get_txs(2).unwrap();

// TODO(Ayelet): all transactions should be returned after replenishing.
// Assert: all remaining transactions are returned.
assert_eq!(txs, &[input_address_0_nonce_0.tx]);
assert_eq!(txs, &[input_address_0_nonce_0.tx, input_address_0_nonce_1.tx]);
}
4 changes: 4 additions & 0 deletions crates/mempool/src/transaction_queue.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,10 @@ impl TransactionQueue {
}
false
}

pub fn is_empty(&self) -> bool {
self.queue.is_empty()
}
}

/// Encapsulates a transaction reference to assess its order (i.e., priority).
Expand Down

0 comments on commit e63abfd

Please sign in to comment.