remove retain_mut and do not retain when polling the handler

Author: jxs

protocols/gossipsub/src/behaviour.rs

@@ -3177,22 +3177,18 @@ where
                     }
                 }
             }
-            HandlerEvent::MessagesDropped(rpcs) => {
+            HandlerEvent::MessagesDropped(rpc) => {
                 // Account for this in the scoring logic
                 if let Some((peer_score, _, _)) = &mut self.peer_score {
                     peer_score.failed_message_slow_peer(&propagation_source);
                 }
 
                 // Keep track of expired messages for the application layer.
                 let failed_messages = self.failed_messages.entry(propagation_source).or_default();
-                failed_messages.timeout += rpcs.len();
+                failed_messages.timeout += 1;
                 if let Some(metrics) = self.metrics.as_mut() {
-                    for rpc in rpcs {
-                        if let RpcOut::Publish { message, .. } | RpcOut::Forward { message, .. } =
-                            rpc
-                        {
-                            metrics.timeout_msg_dropped(&message.topic);
-                        }
+                    if let RpcOut::Publish { message, .. } | RpcOut::Forward { message, .. } = rpc {
+                        metrics.timeout_msg_dropped(&message.topic);
                     }
                 }
             }
@@ -3297,7 +3293,7 @@ where
                             }
 
                             // Remove messages from the queue.
-                            peer.messages.retain_mut(|rpc| match rpc {
+                            peer.messages.retain(|rpc| match rpc {
                                 RpcOut::Publish { message_id, .. }
                                 | RpcOut::Forward { message_id, .. } => {
                                     !message_ids.contains(message_id)

protocols/gossipsub/src/handler.rs

@@ -60,7 +60,7 @@ pub enum HandlerEvent {
     /// which protocol. This message only occurs once per connection.
     PeerKind(PeerKind),
     /// A message to be published was dropped because it could not be sent in time.
-    MessagesDropped(Vec<RpcOut>),
+    MessagesDropped(RpcOut),
 }
 
 /// A message sent from the behaviour to the handler.
@@ -270,7 +270,7 @@ impl EnabledHandler {
                                     self.outbound_substream =
                                         Some(OutboundSubstreamState::WaitingOutput(substream));
                                     return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(
-                                        HandlerEvent::MessagesDropped(vec![message]),
+                                        HandlerEvent::MessagesDropped(message),
                                     ));
                                 }
                             }
@@ -410,22 +410,6 @@ impl EnabledHandler {
             }
         }
 
-        // Remove stale messages from the queue.
-        let stale = self.message_queue.retain_mut(|rpc| match rpc {
-            RpcOut::Publish {
-                ref mut timeout, ..
-            }
-            | RpcOut::Forward {
-                ref mut timeout, ..
-            } => !timeout.poll_unpin(cx).is_ready(),
-            _ => true,
-        });
-        if !stale.is_empty() {
-            return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(
-                HandlerEvent::MessagesDropped(stale),
-            ));
-        }
-
         Poll::Pending
     }
 }

protocols/gossipsub/src/queue.rs

@@ -20,7 +20,6 @@
 
 use std::{
     collections::{BinaryHeap, HashMap},
-    mem,
     sync::{
         atomic::{AtomicUsize, Ordering},
         Arc, Mutex,
@@ -105,21 +104,9 @@ impl<T: Ord> Queue<T> {
     /// Retain only the elements specified by the predicate.
     /// In other words, remove all elements e for which f(&e) returns false. The elements are visited in unsorted (and unspecified) order.
     /// Returns the cleared items.
-    pub(crate) fn retain_mut<F: FnMut(&mut T) -> bool>(&mut self, mut f: F) -> Vec<T> {
+    pub(crate) fn retain<F: FnMut(&T) -> bool>(&mut self, f: F) {
         let mut shared = self.shared.lock().expect("lock to not be poisoned");
-        // `BinaryHeap` doesn't impl `retain_mut`, this seems like a practical way to achieve it.
-        // `BinaryHeap::drain` is O(n) as it returns an iterator over the removed elements in its internal arbitrary order.
-        // `BinaryHeap::push` is ~O(1) which makes this function O(n).
-        let mut queue = mem::replace(&mut shared.queue, BinaryHeap::with_capacity(self.capacity));
-        let mut cleared = vec![];
-        for mut item in queue.drain() {
-            if f(&mut item) {
-                shared.queue.push(item);
-            } else {
-                cleared.push(item);
-            }
-        }
-        cleared
+        shared.queue.retain(f);
     }
 
     /// Returns the length of the queue.