use std::hash::BuildHasher;
use std::collections::{HashSet, HashMap};
use crate::api::TransactionStats;
use vapory_types::{H256, H512};
use fastmap::H256FastMap;
use common_types::BlockNumber;
type NodeId = H512;
#[derive(Debug, PartialEq, Clone, MallocSizeOf)]
pub struct Stats {
first_seen: BlockNumber,
propagated_to: HashMap<NodeId, usize>,
}
impl Stats {
pub fn new(number: BlockNumber) -> Self {
Stats {
first_seen: number,
propagated_to: Default::default(),
}
}
}
impl<'a> From<&'a Stats> for TransactionStats {
fn from(other: &'a Stats) -> Self {
TransactionStats {
first_seen: other.first_seen,
propagated_to: other.propagated_to
.iter()
.map(|(hash, size)| (*hash, *size))
.collect(),
}
}
}
#[derive(Debug, Default, MallocSizeOf)]
pub struct TransactionsStats {
pending_transactions: H256FastMap<Stats>,
}
impl TransactionsStats {
pub fn propagated(&mut self, hash: &H256, enode_id: Option<NodeId>, current_block_num: BlockNumber) {
let enode_id = enode_id.unwrap_or_default();
let stats = self.pending_transactions.entry(*hash).or_insert_with(|| Stats::new(current_block_num));
let count = stats.propagated_to.entry(enode_id).or_insert(0);
*count = count.saturating_add(1);
}
#[cfg(test)]
pub fn get(&self, hash: &H256) -> Option<&Stats> {
self.pending_transactions.get(hash)
}
pub fn stats(&self) -> &H256FastMap<Stats> {
&self.pending_transactions
}
pub fn retain<S: BuildHasher>(&mut self, hashes: &HashSet<H256, S>) {
let to_remove = self.pending_transactions.keys()
.filter(|hash| !hashes.contains(hash))
.cloned()
.collect::<Vec<_>>();
for hash in to_remove {
self.pending_transactions.remove(&hash);
}
}
}
#[cfg(test)]
mod tests {
use std::collections::{HashMap, HashSet};
use super::{Stats, TransactionsStats, NodeId, H256};
use tetsy_macros::hash_map;
#[test]
fn should_keep_track_of_propagations() {
let mut stats = TransactionsStats::default();
let hash = H256::from_low_u64_be(5);
let enodeid1 = NodeId::from_low_u64_be(2);
let enodeid2 = NodeId::from_low_u64_be(5);
stats.propagated(&hash, Some(enodeid1), 5);
stats.propagated(&hash, Some(enodeid1), 10);
stats.propagated(&hash, Some(enodeid2), 15);
let stats = stats.get(&hash);
assert_eq!(stats, Some(&Stats {
first_seen: 5,
propagated_to: hash_map![
enodeid1 => 2,
enodeid2 => 1
],
}));
}
#[test]
fn should_remove_hash_from_tracking() {
let mut stats = TransactionsStats::default();
let hash = H256::from_low_u64_be(5);
let enodeid1 = NodeId::from_low_u64_be(5);
stats.propagated(&hash, Some(enodeid1), 10);
stats.retain(&HashSet::new());
let stats = stats.get(&hash);
assert_eq!(stats, None);
}
}