feat: reset route optimization

informalsystems · Jan 10, 2025 · f5b1fc8 · f5b1fc8
1 parent c8f429f
commit f5b1fc8
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 9 deletions.
diff --git a/dog/src/behaviour.rs b/dog/src/behaviour.rs
@@ -16,7 +16,6 @@ use libp2p::{
 };
 use lru::LruCache;
 use quick_protobuf::{MessageWrite, Writer};
-use rand::seq::IteratorRandom;
 
 use crate::{
     config::Config,
@@ -447,14 +446,16 @@ where
 
             tracing::debug!(peer=%propagation_source, "Sending HaveTx to peer");
 
-            self.send_transaction(
+            if self.send_transaction(
                 *propagation_source,
                 RpcOut::HaveTx(HaveTx {
                     from: transaction.from,
                 }),
-            );
+            ) {
+                self.router.register_have_tx_sent(*propagation_source);
+                self.redundancy_controller.block_have_tx();
+            }
 
-            self.redundancy_controller.block_have_tx();
             return;
         }
         self.redundancy_controller.incr_first_time_txs_count();
@@ -516,11 +517,12 @@ where
         if self.redundancy_controller.evaluate() {
             tracing::warn!("Redundancy is too low. Sending reset route");
 
-            let mut rng = rand::thread_rng();
-            match self.connected_peers.keys().choose(&mut rng) {
+            match self.router.get_random_have_tx_sent_peer() {
                 Some(peer_id) => {
                     tracing::trace!(peer=%peer_id, "Sending reset route to peer");
-                    self.send_transaction(*peer_id, RpcOut::ResetRoute(ResetRoute {}));
+                    if self.send_transaction(peer_id, RpcOut::ResetRoute(ResetRoute {})) {
+                        self.router.remove_have_tx_sent(&peer_id);
+                    }
                 }
                 None => {
                     // This should not happen

diff --git a/dog/src/dog.rs b/dog/src/dog.rs
@@ -1,4 +1,4 @@
-use std::fmt::Display;
+use std::{collections::HashMap, fmt::Display};
 
 use libp2p::PeerId;
 use rand::seq::IteratorRandom;
@@ -32,14 +32,15 @@ impl Display for Route {
 }
 
 pub(crate) struct Router {
-    // TODO: is it better to use a HashMap<PeerId, HashSet<PeerId>>?
     disabled_routes: Vec<Route>,
+    have_tx_sent_per_peer: HashMap<PeerId, usize>,
 }
 
 impl Router {
     pub(crate) fn new() -> Self {
         Router {
             disabled_routes: Vec::new(),
+            have_tx_sent_per_peer: HashMap::new(),
         }
     }
 
@@ -88,6 +89,27 @@ impl Router {
             })
             .collect()
     }
+
+    pub(crate) fn register_have_tx_sent(&mut self, peer: PeerId) {
+        let counter = self.have_tx_sent_per_peer.entry(peer).or_insert(0);
+        *counter += 1;
+    }
+
+    // Returns a random peer to which we have sent a have_tx message.
+    pub(crate) fn get_random_have_tx_sent_peer(&self) -> Option<PeerId> {
+        self.have_tx_sent_per_peer
+            .iter()
+            .filter_map(|(peer, count)| if *count > 0 { Some(peer.clone()) } else { None })
+            .choose(&mut rand::thread_rng())
+    }
+
+    pub(crate) fn remove_have_tx_sent(&mut self, peer: &PeerId) {
+        if let Some(counter) = self.have_tx_sent_per_peer.get_mut(peer) {
+            if *counter > 0 {
+                *counter -= 1;
+            }
+        }
+    }
 }
 
 pub(crate) struct Controller {