diff --git a/crates/mempool/src/lib.rs b/crates/mempool/src/lib.rs index 6d5d6f00..a43dfc00 100644 --- a/crates/mempool/src/lib.rs +++ b/crates/mempool/src/lib.rs @@ -1,3 +1,4 @@ pub mod mempool; pub(crate) mod priority_queue; +pub mod staging_area; pub mod transaction_store; diff --git a/crates/mempool/src/mempool.rs b/crates/mempool/src/mempool.rs index 3112a6b9..be344755 100644 --- a/crates/mempool/src/mempool.rs +++ b/crates/mempool/src/mempool.rs @@ -13,6 +13,7 @@ use starknet_mempool_types::mempool_types::{ use tokio::sync::mpsc::Receiver; use crate::priority_queue::TransactionPriorityQueue; +use crate::staging_area::StagingQueue; use crate::transaction_store::TransactionStore; #[cfg(test)] @@ -23,6 +24,7 @@ pub struct Mempool { // TODO: add docstring explaining visibility and coupling of the fields. txs_queue: TransactionPriorityQueue, tx_store: TransactionStore, + staging: StagingQueue, state: HashMap, } @@ -32,6 +34,7 @@ impl Mempool { txs_queue: TransactionPriorityQueue::default(), tx_store: TransactionStore::default(), state: HashMap::default(), + staging: StagingQueue::default(), }; for input in inputs.into_iter() { @@ -77,6 +80,7 @@ impl Mempool { for pq_tx in &pq_txs { let tx = self.tx_store.remove(&pq_tx.tx_hash)?; self.state.remove(&tx.sender_address); + self.staging.insert(tx.clone().into())?; txs.push(tx); } diff --git a/crates/mempool/src/staging_area.rs b/crates/mempool/src/staging_area.rs new file mode 100644 index 00000000..ab9b4fea --- /dev/null +++ b/crates/mempool/src/staging_area.rs @@ -0,0 +1,63 @@ +use std::cmp::Ordering; +use std::collections::HashMap; + +use starknet_api::core::{ContractAddress, Nonce}; +use starknet_api::transaction::TransactionHash; +use starknet_mempool_types::errors::MempoolError; +use starknet_mempool_types::mempool_types::ThinTransaction; + +#[derive(Clone, Debug, Default)] +pub struct StagingTransaction { + pub tx_hash: TransactionHash, + pub address: ContractAddress, + pub nonce: Nonce, +} + +impl PartialEq for StagingTransaction { + fn eq(&self, other: &StagingTransaction) -> bool { + self.address == other.address && self.nonce == other.nonce + } +} + +impl Eq for StagingTransaction {} + +impl Ord for StagingTransaction { + fn cmp(&self, other: &Self) -> Ordering { + self.address.cmp(&other.address).then_with(|| self.nonce.cmp(&other.nonce)) + } +} + +impl PartialOrd for StagingTransaction { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +#[derive(Clone, Debug, Default, derive_more::Deref, derive_more::DerefMut)] +pub struct StagingQueue(HashMap); +impl StagingQueue { + pub fn insert(&mut self, tx: StagingTransaction) -> Result<(), MempoolError> { + if self.0.contains_key(&tx.tx_hash) { + return Err(MempoolError::DuplicateTransaction { tx_hash: tx.tx_hash }); + } + self.0.insert(tx.tx_hash, tx); + + Ok(()) + } + + pub fn remove( + &mut self, + tx_hash: &TransactionHash, + ) -> Result { + match self.0.remove(tx_hash) { + Some(tx) => Ok(tx), + None => Err(MempoolError::TransactionNotFound { tx_hash: *tx_hash }), + } + } +} + +impl From for StagingTransaction { + fn from(tx: ThinTransaction) -> Self { + StagingTransaction { address: tx.sender_address, nonce: tx.nonce, tx_hash: tx.tx_hash } + } +}