nu-node/crates/nu-mempool/src/pool.rs

use std::collections::HashMap;

use nu_block::types::RawTransaction;

use crate::priority::TxPriority;
use crate::types::{MEMPOOL_MAX_TX, MEMPOOL_TTL_MS, MAX_TX_PER_SENDER};

#[derive(Debug, Clone)]
pub struct PendingTx {
    pub tx:          RawTransaction,
    pub priority:    TxPriority,
    pub received_at: i64, // Unix epoch ms
}

pub struct Mempool {
    txs:           Vec<PendingTx>,
    sender_counts: HashMap<String, usize>,
}

impl Mempool {
    pub fn new() -> Self {
        Self { txs: vec![], sender_counts: HashMap::new() }
    }

    pub fn insert(&mut self, tx: RawTransaction, now_ms: i64) -> bool {
        if self.txs.len() >= MEMPOOL_MAX_TX {
            return false;
        }
        let priority = TxPriority::from_payload(&tx.payload);
        let count = self.sender_counts.entry(tx.sender.clone()).or_insert(0);
        if *count >= MAX_TX_PER_SENDER && priority == TxPriority::Normal {
            return false;
        }
        *count += 1;
        self.txs.push(PendingTx { tx, priority, received_at: now_ms });
        true
    }

    pub fn contains(&self, tx_id: &str) -> bool {
        self.txs.iter().any(|p| p.tx.tx_id == tx_id)
    }

    pub fn select_for_block(&mut self, max_tx: usize, now_ms: i64) -> Vec<PendingTx> {
        self.evict_expired(now_ms);
        let mut sorted = self.txs.clone();
        sorted.sort_by(|a, b| {
            b.priority.cmp(&a.priority).then(b.tx.fee.cmp(&a.tx.fee))
        });
        sorted.into_iter().take(max_tx).collect()
    }

    pub fn remove(&mut self, tx_id: &str) {
        if let Some(pos) = self.txs.iter().position(|p| p.tx.tx_id == tx_id) {
            let removed = self.txs.remove(pos);
            if let Some(c) = self.sender_counts.get_mut(&removed.tx.sender) {
                *c = c.saturating_sub(1);
            }
        }
    }

    pub fn len(&self) -> usize {
        self.txs.len()
    }

    fn evict_expired(&mut self, now_ms: i64) {
        let mut removed_senders: Vec<String> = vec![];
        self.txs.retain(|p| {
            let keep = now_ms - p.received_at < MEMPOOL_TTL_MS;
            if !keep {
                removed_senders.push(p.tx.sender.clone());
            }
            keep
        });
        for sender in removed_senders {
            if let Some(c) = self.sender_counts.get_mut(&sender) {
                *c = c.saturating_sub(1);
            }
        }
    }
}