From c8f248bd896b33f83c53a82253bb8243224c5e5c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Oliveira?= <hello@jxs.pt>
Date: Tue, 5 Dec 2023 19:35:18 +0000
Subject: [PATCH] send IHAVE messages immediately when emitting gossip

---
 protocols/gossipsub/src/behaviour.rs       | 19 ++++++++-----------
 protocols/gossipsub/src/behaviour/tests.rs | 20 ++++++++++++--------
 2 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/protocols/gossipsub/src/behaviour.rs b/protocols/gossipsub/src/behaviour.rs
index 0108a645c17f..19d5d9e728c7 100644
--- a/protocols/gossipsub/src/behaviour.rs
+++ b/protocols/gossipsub/src/behaviour.rs
@@ -2516,14 +2516,14 @@ where
                 }
 
                 // send an IHAVE message
-                Self::control_pool_add(
-                    &mut self.control_pool,
-                    peer,
-                    RpcOut::IHave(IHave {
-                        topic_hash: topic_hash.clone(),
-                        message_ids: peer_message_ids,
-                    }),
-                );
+                let sender = self
+                    .handler_send_queues
+                    .get_mut(&peer)
+                    .expect("Peerid should exist");
+                sender.ihave(IHave {
+                    topic_hash: topic_hash.clone(),
+                    message_ids: peer_message_ids,
+                });
             }
         }
     }
@@ -2808,9 +2808,6 @@ where
 
                 match msg {
                     RpcOut::IHave(ihave) => sender.ihave(ihave),
-                    RpcOut::IWant(iwant) => sender.iwant(iwant),
-                    RpcOut::Graft(graft) => sender.graft(graft),
-                    RpcOut::Prune(prune) => sender.prune(prune),
                     _ => unreachable!(),
                 }
             }
diff --git a/protocols/gossipsub/src/behaviour/tests.rs b/protocols/gossipsub/src/behaviour/tests.rs
index 9d2be223d032..b0a61f4510f3 100644
--- a/protocols/gossipsub/src/behaviour/tests.rs
+++ b/protocols/gossipsub/src/behaviour/tests.rs
@@ -2538,7 +2538,7 @@ fn test_do_not_gossip_to_peers_below_gossip_threshold() {
     };
 
     // Build full mesh
-    let (mut gs, peers, queues, topics) = inject_nodes1()
+    let (mut gs, peers, mut receivers, topics) = inject_nodes1()
         .peer_no(config.mesh_n_high())
         .topics(vec!["test".into()])
         .to_subscribe(true)
@@ -2552,8 +2552,10 @@ fn test_do_not_gossip_to_peers_below_gossip_threshold() {
     }
 
     // Add two additional peers that will not be part of the mesh
-    let (p1, _receiver1) = add_peer(&mut gs, &topics, false, false);
-    let (p2, _receiver2) = add_peer(&mut gs, &topics, false, false);
+    let (p1, receiver1) = add_peer(&mut gs, &topics, false, false);
+    receivers.insert(p1, receiver1);
+    let (p2, receiver2) = add_peer(&mut gs, &topics, false, false);
+    receivers.insert(p2, receiver2);
 
     // Reduce score of p1 below peer_score_thresholds.gossip_threshold
     // note that penalties get squared so two penalties means a score of
@@ -2585,7 +2587,7 @@ fn test_do_not_gossip_to_peers_below_gossip_threshold() {
 
     // Check that exactly one gossip messages got sent and it got sent to p2
     assert_eq!(
-        count_control_msgs(&gs, &queues, |peer, action| match action {
+        count_control_msgs(&gs, &receivers, |peer, action| match action {
             RpcOut::IHave(IHave {
                 topic_hash,
                 message_ids,
@@ -4568,7 +4570,7 @@ fn test_limit_number_of_message_ids_inside_ihave() {
         .build()
         .unwrap();
     //build gossipsub with full mesh
-    let (mut gs, peers, queues, topics) = inject_nodes1()
+    let (mut gs, peers, mut receivers, topics) = inject_nodes1()
         .peer_no(config.mesh_n_high())
         .topics(vec!["test".into()])
         .to_subscribe(false)
@@ -4581,8 +4583,10 @@ fn test_limit_number_of_message_ids_inside_ihave() {
     }
 
     //add two other peers not in the mesh
-    let (p1, _) = add_peer(&mut gs, &topics, false, false);
-    let (p2, _) = add_peer(&mut gs, &topics, false, false);
+    let (p1, receiver1) = add_peer(&mut gs, &topics, false, false);
+    receivers.insert(p1, receiver1);
+    let (p2, receiver2) = add_peer(&mut gs, &topics, false, false);
+    receivers.insert(p2, receiver2);
 
     //receive 200 messages from another peer
     let mut seq = 0;
@@ -4601,7 +4605,7 @@ fn test_limit_number_of_message_ids_inside_ihave() {
     let mut ihaves2 = HashSet::new();
 
     assert_eq!(
-        count_control_msgs(&gs, &queues, |p, action| match action {
+        count_control_msgs(&gs, &receivers, |p, action| match action {
             RpcOut::IHave(IHave { message_ids, .. }) => {
                 if p == &p1 {
                     ihaves1 = message_ids.iter().cloned().collect();