Skip to content

Commit

Permalink
feat: add staging area
Browse files Browse the repository at this point in the history
  • Loading branch information
MohammadNassar1 committed Jun 4, 2024
1 parent 03ce42f commit a3b6b1e
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 0 deletions.
1 change: 1 addition & 0 deletions crates/mempool/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
pub mod mempool;
pub(crate) mod priority_queue;
pub mod staging_area;
pub mod transaction_store;
4 changes: 4 additions & 0 deletions crates/mempool/src/mempool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -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<ContractAddress, AccountState>,
}

Expand All @@ -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() {
Expand Down Expand Up @@ -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);
}

Expand Down
51 changes: 51 additions & 0 deletions crates/mempool/src/staging_area.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
use std::cmp::Ordering;
use std::collections::HashMap;

use starknet_api::core::{ContractAddress, Nonce};
use starknet_api::transaction::TransactionHash;
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<Ordering> {
Some(self.cmp(other))
}
}

#[derive(Clone, Debug, Default, derive_more::Deref, derive_more::DerefMut)]
pub struct StagingQueue(HashMap<TransactionHash, StagingTransaction>);
impl StagingQueue {
pub fn insert(&mut self, tx: StagingTransaction) {
self.0.insert(tx.tx_hash, tx);
}

pub fn remove(&mut self, tx_hash: &TransactionHash) -> Option<StagingTransaction> {
self.0.remove(tx_hash)
}
}

impl From<ThinTransaction> for StagingTransaction {
fn from(tx: ThinTransaction) -> Self {
StagingTransaction { address: tx.sender_address, nonce: tx.nonce, tx_hash: tx.tx_hash }
}
}

0 comments on commit a3b6b1e

Please sign in to comment.