feat(wallet): WIP: BroadcastQueue

Author: evanlinjin

wallet/src/wallet/broadcast_queue.rs

@@ -0,0 +1,196 @@
+//! Unbroadcasted transaction queue.
+
+use std::vec::Vec;
+
+use crate::collections::BTreeMap;
+use crate::collections::HashMap;
+use crate::collections::VecDeque;
+
+use bitcoin::Txid;
+use chain::Merge;
+use chain::TxUpdate;
+
+/// An ordered unbroadcasted list.
+///
+/// It is ordered in case of RBF txs.
+#[derive(Debug, Clone, Default)]
+pub struct BroadcastQueue {
+    queue: VecDeque<Txid>,
+
+    /// Enforces that we only change the tx entry if the change has a larger `seq`.
+    latest_seq: HashMap<Txid, (u64, bool)>,
+    next_seq: u64,
+}
+
+/// Represents changes made to [`BroadcastQueue`].
+#[must_use]
+#[derive(Debug, Clone, Default, PartialEq, serde::Deserialize, serde::Serialize)]
+pub struct ChangeSet {
+    pub insert_tx: BTreeMap<Txid, u64>,
+    pub remove_tx: BTreeMap<Txid, u64>,
+}
+
+impl Merge for ChangeSet {
+    fn merge(&mut self, other: Self) {
+        self.insert_tx.extend(
+            other
+                .insert_tx
+                .into_iter()
+                // Only replace `seq`uences which are larger.
+                .filter(|(txid, other_seq)| match self.insert_tx.get(txid) {
+                    Some(seq) => seq < other_seq,
+                    None => true,
+                })
+                .collect::<Vec<_>>(),
+        );
+        self.remove_tx.extend(
+            other
+                .remove_tx
+                .into_iter()
+                .filter(|(txid, other_seq)| match self.remove_tx.get(txid) {
+                    Some(seq) => seq < other_seq,
+                    None => true,
+                })
+                .collect::<Vec<_>>(),
+        );
+    }
+
+    fn is_empty(&self) -> bool {
+        self.insert_tx.is_empty() && self.remove_tx.is_empty()
+    }
+}
+
+impl BroadcastQueue {
+    /// Construct [`Unbroadcasted`] from the given `changeset`.
+    pub fn from_changeset(changeset: ChangeSet) -> Self {
+        let mut out = BroadcastQueue::default();
+        out.apply_changeset(changeset);
+        out
+    }
+
+    /// Apply the given `changeset`.
+    pub fn apply_changeset(&mut self, changeset: ChangeSet) {
+        for (txid, seq) in changeset.insert_tx {
+            self._push(txid, seq);
+        }
+        for (txid, seq) in changeset.remove_tx {
+            self._remove(txid, seq);
+        }
+        self._sort_queue();
+    }
+
+    fn _sort_queue(&mut self) {
+        self.queue
+            .make_contiguous()
+            .sort_unstable_by_key(|txid| self.latest_seq.get(txid));
+    }
+
+    fn _push(&mut self, txid: Txid, seq: u64) -> bool {
+        let modified = &mut false;
+        self.latest_seq
+            .entry(txid)
+            .and_modify(|(v, in_queue)| {
+                if *v < seq {
+                    *v = seq;
+                    *modified = true;
+                    if !*in_queue {
+                        *in_queue = true;
+                        self.queue.push_back(txid);
+                    }
+                }
+            })
+            .or_insert_with(|| {
+                *modified = true;
+                self.queue.push_back(txid);
+                (seq, true)
+            });
+        self.next_seq = self.next_seq.max(seq.saturating_add(1));
+        *modified
+    }
+
+    fn _contains(latest_seq: &mut HashMap<Txid, (u64, bool)>, txid: Txid) -> bool {
+        match latest_seq.get(&txid) {
+            Some(&(_, contains)) => contains,
+            None => false,
+        }
+    }
+
+    fn _remove(&mut self, txid: Txid, seq: u64) -> bool {
+        let modified = &mut false;
+        self.latest_seq
+            .entry(txid)
+            .and_modify(|(v, in_queue)| {
+                if *v <= seq {
+                    *v = seq;
+                    *modified = true;
+                    if *in_queue {
+                        *in_queue = false;
+                        // Most of the time, we are popping from the front of the queue.
+                        let index = self
+                            .queue
+                            .iter()
+                            .enumerate()
+                            .find_map(|(i, &q_txid)| if q_txid == txid { Some(i) } else { None })
+                            .expect("must be in queue");
+                        self.queue.swap_remove_front(index);
+                    }
+                }
+            })
+            .or_insert_with(|| {
+                *modified = true;
+                (seq, false)
+            });
+        self.next_seq = self.next_seq.max(seq.saturating_add(1));
+        *modified
+    }
+
+    /// Reinserting will bump the tx's seq to `next_seq`.
+    pub fn push(&mut self, txid: Txid) -> ChangeSet {
+        let seq = self.next_seq;
+
+        let mut changeset = ChangeSet::default();
+        if self._push(txid, seq) {
+            changeset.insert_tx.insert(txid, seq);
+            self.next_seq += 1;
+        }
+        changeset
+    }
+
+    /// Returns the next `txid` of the transaction to broadcast.
+    pub fn next(&self) -> Option<Txid> {
+        // TODO: Implement.
+        // NOTE: Prioritize txs that have no dependencies.
+        None
+    }
+
+    /// Untrack the `txid`.
+    pub fn remove(&mut self, txid: Txid) -> ChangeSet {
+        let seq = self.next_seq;
+
+        let mut changeset = ChangeSet::default();
+        if self._remove(txid, seq) {
+            changeset.remove_tx.insert(txid, seq);
+            self.next_seq += 1;
+        }
+        changeset
+    }
+
+    /// Untrack transactions that are given anchors and seen-at timestamps.
+    pub fn update<A>(&mut self, tx_update: &TxUpdate<A>) -> ChangeSet {
+        let mut changeset = ChangeSet::default();
+        for (_, txid) in &tx_update.anchors {
+            changeset.merge(self.remove(*txid));
+        }
+        for (txid, _) in &tx_update.seen_ats {
+            changeset.merge(self.remove(*txid));
+        }
+        changeset
+    }
+
+    /// Txids ordered by precedence.
+    ///
+    /// Transactions with greater precedence will appear later in this list.
+    pub fn txids(&self) -> impl ExactSizeIterator<Item = Txid> + '_ {
+        self.queue.iter().copied()
+    }
+}

wallet/src/wallet/mod.rs

@@ -58,6 +58,7 @@ mod persisted;
 pub mod signer;
 pub mod tx_builder;
 pub(crate) mod utils;
+mod broadcast_queue;
 
 use crate::collections::{BTreeMap, HashMap, HashSet};
 use crate::descriptor::{
@@ -82,6 +83,7 @@ pub use params::*;
 pub use persisted::*;
 pub use utils::IsDust;
 pub use utils::TxDetails;
+pub use broadcast_queue::*;
 
 /// A Bitcoin wallet
 ///