Refactor mempool tracker to use HashSet difference

src/bin/tracker.rs

@@ -16,10 +16,10 @@ fn main() {
     ]).unwrap();
 
     let daemon = daemon::Daemon::new("localhost:8332");
-    let mut tracker = mempool::Tracker::new(&daemon);
+    let mut tracker = mempool::Tracker::new();
     loop {
         let t = Instant::now();
-        tracker.update_from_daemon().unwrap();
+        tracker.update(&daemon).unwrap();
         let dt = t.elapsed();
         info!(
             "update took {:.3} ms",

src/daemon.rs

@@ -38,11 +38,20 @@ pub struct Daemon {
 pub struct MempoolEntry {
     fee: u64,   // in satoshis
     vsize: u32, // in virtual bytes (= weight/4)
+    fee_per_vbyte: f32,
 }
 
 impl MempoolEntry {
+    fn new(fee: u64, vsize: u32) -> MempoolEntry {
+        MempoolEntry {
+            fee,
+            vsize,
+            fee_per_vbyte: fee as f32 / vsize as f32,
+        }
+    }
+
     pub fn fee_per_vbyte(&self) -> f32 {
-        self.fee as f32 / self.vsize as f32
+        self.fee_per_vbyte
     }
 
     pub fn fee(&self) -> u64 {
@@ -201,17 +210,16 @@ impl Daemon {
             .chain_err(|| "missing fees section")?
             .as_object()
             .chain_err(|| "non-object fees")?;
-        Ok(MempoolEntry {
-            fee: (fees.get("base")
+        let fee = (fees.get("base")
             .chain_err(|| "missing base fee")?
             .as_f64()
-                .chain_err(|| "non-float fee")? * 100_000_000f64) as u64,
-            vsize: entry
+            .chain_err(|| "non-float fee")? * 100_000_000f64) as u64;
+        let vsize = entry
             .get("size")
             .chain_err(|| "missing size")?
             .as_u64()
-                .chain_err(|| "non-integer size")? as u32,
-        })
+            .chain_err(|| "non-integer size")? as u32;
+        Ok(MempoolEntry::new(fee, vsize))
     }
 
     pub fn get_all_headers(&self) -> Result<HeaderMap> {

src/mempool.rs

@@ -1,7 +1,6 @@
 use bitcoin::blockdata::transaction::Transaction;
 use bitcoin::util::hash::Sha256dHash;
 use daemon::{Daemon, MempoolEntry};
-use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet};
 use std::iter::FromIterator;
 
@@ -14,16 +13,20 @@ pub struct Stats {
     entry: MempoolEntry,
 }
 
-pub struct Tracker<'a> {
-    txids: HashMap<Sha256dHash, Stats>,
-    daemon: &'a Daemon,
+impl Stats {
+    pub fn new(tx: Transaction, entry: MempoolEntry) -> Stats {
+        Stats { tx, entry }
+    }
 }
 
-impl<'a> Tracker<'a> {
-    pub fn new(daemon: &Daemon) -> Tracker {
+pub struct Tracker {
+    stats: HashMap<Sha256dHash, Stats>,
+}
+
+impl Tracker {
+    pub fn new() -> Tracker {
         Tracker {
-            txids: HashMap::new(),
-            daemon: daemon,
+            stats: HashMap::new(),
         }
     }
 
@@ -31,7 +34,7 @@ impl<'a> Tracker<'a> {
     /// cumulative virtual size of mempool transaction with fee in the interval [fee_{n-1}, fee_n].
     /// Note: fee_0 is implied to be infinity.
     pub fn fee_histogram(&self) -> Vec<(f32, u32)> {
-        let mut entries: Vec<&MempoolEntry> = self.txids.values().map(|stat| &stat.entry).collect();
+        let mut entries: Vec<&MempoolEntry> = self.stats.values().map(|stat| &stat.entry).collect();
         entries.sort_unstable_by(|e1, e2| {
             e2.fee_per_vbyte().partial_cmp(&e1.fee_per_vbyte()).unwrap()
         });
@@ -48,21 +51,22 @@ impl<'a> Tracker<'a> {
         histogram
     }
 
-    pub fn update_from_daemon(&mut self) -> Result<()> {
-        let new_txids = HashSet::<Sha256dHash>::from_iter(self.daemon
+    pub fn update(&mut self, daemon: &Daemon) -> Result<()> {
+        let new_txids = HashSet::<Sha256dHash>::from_iter(daemon
             .getmempooltxids()
             .chain_err(|| "failed to update mempool from daemon")?);
-        self.txids.retain(|txid, _| new_txids.contains(txid)); // remove old TXNs
-        for txid in new_txids {
-            if let Entry::Vacant(map_entry) = self.txids.entry(txid) {
-                let tx = match self.daemon.gettransaction(&txid) {
+        let old_txids = HashSet::from_iter(self.stats.keys().cloned());
+
+        let mut to_add = Vec::new();
+        for &txid in new_txids.difference(&old_txids) {
+            let tx = match daemon.gettransaction(&txid) {
                 Ok(tx) => tx,
                 Err(err) => {
                     warn!("missing tx {}: {}", txid, err);
                     continue;
                 }
             };
-                let entry = match self.daemon.getmempoolentry(&txid) {
+            let entry = match daemon.getmempoolentry(&txid) {
                 Ok(entry) => entry,
                 Err(err) => {
                     warn!("no mempool entry {}: {}", txid, err);
@@ -70,9 +74,13 @@ impl<'a> Tracker<'a> {
                 }
             };
             debug!("new tx: {}, {:.3}", txid, entry.fee_per_vbyte(),);
-                map_entry.insert(Stats { tx, entry });
+            to_add.push((txid, Stats::new(tx, entry)));
         }
+        self.stats.extend(to_add);
+        for txid in old_txids.difference(&new_txids) {
+            self.stats.remove(txid);
         }
+        assert_eq!(new_txids, HashSet::from_iter(self.stats.keys().cloned()));
         Ok(())
     }
 }