Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: update priority queue to use ThinPriorityTransaction #195

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 21 additions & 13 deletions crates/mempool/src/mempool.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::collections::HashMap;

use starknet_api::core::ContractAddress;
use starknet_api::transaction::TransactionHash;
use starknet_api::core::{ContractAddress, Nonce};
use starknet_api::transaction::{Tip, TransactionHash};
use starknet_mempool_types::mempool_types::{
AccountState, MempoolInput, MempoolResult, ThinTransaction,
};
Expand Down Expand Up @@ -38,7 +38,7 @@ impl Mempool {

/// Returns an iterator of the current eligible transactions for sequencing, ordered by their
/// priority.
pub fn iter(&self) -> impl Iterator<Item = TransactionReference> + '_ {
pub fn iter(&self) -> impl Iterator<Item = &TransactionReference> {
self.tx_queue.iter()
}

Expand All @@ -49,11 +49,9 @@ 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);

let txs = self.tx_queue.pop_last_chunk(n_txs);
for tx in txs {
self.tx_pool.remove(tx.tx_hash)?;
eligible_txs.push(tx.0);
for tx_hash in self.tx_queue.pop_last_chunk(n_txs) {
let tx = self.tx_pool.remove(tx_hash)?;
eligible_txs.push(tx);
}

Ok(eligible_txs)
Expand Down Expand Up @@ -83,7 +81,7 @@ impl Mempool {
let tx = input.tx;

self.tx_pool.insert(tx.clone())?;
self.tx_queue.insert(TransactionReference::new(tx));
self.tx_queue.insert(TransactionReference::new(&tx));

Ok(())
}
Expand All @@ -93,11 +91,21 @@ impl Mempool {
/// execution fields).
/// TODO(Mohammad): rename this struct to `ThinTransaction` once that name
/// becomes available, to better reflect its purpose and usage.
#[derive(Clone, Debug, Default, derive_more::Deref)]
pub struct TransactionReference(pub ThinTransaction);
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct TransactionReference {
pub sender_address: ContractAddress,
pub nonce: Nonce,
pub tx_hash: TransactionHash,
pub tip: Tip,
}

impl TransactionReference {
pub fn new(tx: ThinTransaction) -> Self {
TransactionReference(tx)
pub fn new(tx: &ThinTransaction) -> Self {
TransactionReference {
sender_address: tx.sender_address,
nonce: tx.nonce,
tx_hash: tx.tx_hash,
tip: tx.tip,
}
}
}
5 changes: 3 additions & 2 deletions crates/mempool/src/mempool_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use starknet_api::{contract_address, patricia_key};
use starknet_mempool_types::errors::MempoolError;
use starknet_mempool_types::mempool_types::{Account, ThinTransaction};

use crate::mempool::{Mempool, MempoolInput};
use crate::mempool::{Mempool, MempoolInput, TransactionReference};

/// Creates a valid input for mempool's `add_tx` with optional default value for
/// `sender_address`.
Expand Down Expand Up @@ -49,11 +49,12 @@ fn mempool() -> Mempool {
#[track_caller]
fn check_mempool_txs_eq(mempool: &Mempool, expected_txs: &[ThinTransaction]) {
let mempool_txs = mempool.tx_queue.iter();
let expected_txs = expected_txs.iter().map(TransactionReference::new);

assert!(
zip_eq(expected_txs, mempool_txs)
// Deref the inner mempool tx type.
.all(|(expected_tx, mempool_tx)| *expected_tx == *mempool_tx)
.all(|(expected_tx, mempool_tx)| expected_tx == *mempool_tx)
);
}

Expand Down
14 changes: 8 additions & 6 deletions crates/mempool/src/transaction_pool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,22 +20,24 @@ pub struct TransactionPool {
}

impl TransactionPool {
// TODO(Mohammad): Remove the cloning of tx once the `TransactionReference` is updated.
pub fn insert(&mut self, tx: ThinTransaction) -> MempoolResult<()> {
let tx_hash = tx.tx_hash;
let tx_reference = TransactionReference::new(&tx);
let tx_hash = tx_reference.tx_hash;

// Insert transaction to pool, if it is new.
if let hash_map::Entry::Vacant(entry) = self.tx_pool.entry(tx_hash) {
entry.insert(tx.clone());
entry.insert(tx);
} else {
return Err(MempoolError::DuplicateTransaction { tx_hash });
}

let txs_from_account_entry = self.txs_by_account.entry(tx.sender_address).or_default();
match txs_from_account_entry.entry(tx.nonce) {
let txs_from_account_entry =
self.txs_by_account.entry(tx_reference.sender_address).or_default();
match txs_from_account_entry.entry(tx_reference.nonce) {
btree_map::Entry::Vacant(txs_from_account) => {
txs_from_account.insert(TransactionReference::new(tx));
txs_from_account.insert(tx_reference);
}
// TODO: support fee escalation transactions.
btree_map::Entry::Occupied(_) => {
panic!(
"Transaction pool consistency error: transaction with hash {tx_hash} does not \
Expand Down
10 changes: 5 additions & 5 deletions crates/mempool/src/transaction_queue.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::cmp::Ordering;
use std::collections::{BTreeSet, HashMap};

use starknet_api::core::{ContractAddress, Nonce};
use starknet_api::transaction::TransactionHash;

use crate::mempool::TransactionReference;
// Assumption: for the MVP only one transaction from the same contract class can be in the mempool
Expand All @@ -28,19 +29,18 @@ impl TransactionQueue {
}

// TODO(gilad): remove collect
pub fn pop_last_chunk(&mut self, n_txs: usize) -> Vec<TransactionReference> {
pub fn pop_last_chunk(&mut self, n_txs: usize) -> Vec<TransactionHash> {
let txs: Vec<TransactionReference> =
(0..n_txs).filter_map(|_| self.queue.pop_last().map(|tx| tx.0)).collect();

for tx in &txs {
self.address_to_nonce.remove(&tx.sender_address);
}

txs
txs.into_iter().map(|tx| tx.tx_hash).collect()
}

pub fn iter(&self) -> impl Iterator<Item = TransactionReference> + '_ {
self.queue.iter().map(|tx| tx.0.clone())
pub fn iter(&self) -> impl Iterator<Item = &TransactionReference> {
self.queue.iter().map(|tx| &tx.0)
}

pub fn _get_nonce(&self, address: &ContractAddress) -> Option<&Nonce> {
Expand Down
Loading