apollo_network: fixed memory hoarding in streams

Author: sirandreww-starkware

crates/apollo_network/src/sqmr/behaviour.rs

@@ -112,7 +112,7 @@ pub struct Behaviour {
     next_outbound_session_id: OutboundSessionId,
     next_inbound_session_id: Arc<AtomicUsize>,
     dropped_sessions: HashSet<SessionId>,
-    wakers_waiting_for_event: Vec<Waker>,
+    last_waker_waiting_for_event: Option<Waker>,
     outbound_sessions_pending_peer_assignment: HashMap<OutboundSessionId, (Bytes, StreamProtocol)>,
     supported_inbound_protocols: HashSet<StreamProtocol>,
 }
@@ -126,7 +126,7 @@ impl Behaviour {
             next_outbound_session_id: Default::default(),
             next_inbound_session_id: Arc::new(Default::default()),
             dropped_sessions: Default::default(),
-            wakers_waiting_for_event: Default::default(),
+            last_waker_waiting_for_event: None,
             outbound_sessions_pending_peer_assignment: Default::default(),
             supported_inbound_protocols: Default::default(),
         }
@@ -213,7 +213,7 @@ impl Behaviour {
 
     fn add_event_to_queue(&mut self, event: ToSwarm<Event, RequestFromBehaviourEvent>) {
         self.pending_events.push_back(event);
-        for waker in self.wakers_waiting_for_event.drain(..) {
+        if let Some(waker) = self.last_waker_waiting_for_event.take() {
             waker.wake();
         }
     }
@@ -330,7 +330,7 @@ impl NetworkBehaviour for Behaviour {
         if let Some(event) = self.pending_events.pop_front() {
             return Poll::Ready(event);
         }
-        self.wakers_waiting_for_event.push(cx.waker().clone());
+        self.last_waker_waiting_for_event = Some(cx.waker().clone());
         Poll::Pending
     }
 }

crates/apollo_network/src/sqmr/handler/inbound_session.rs

@@ -16,7 +16,7 @@ use crate::Bytes;
 pub(super) struct InboundSession {
     pending_messages: VecDeque<Bytes>,
     current_task: WriteMessageTask,
-    wakers_waiting_for_new_message: Vec<Waker>,
+    last_waker_waiting_for_new_message: Option<Waker>,
 }
 
 enum FinishReason {
@@ -35,13 +35,13 @@ impl InboundSession {
         Self {
             pending_messages: Default::default(),
             current_task: WriteMessageTask::Waiting(write_stream),
-            wakers_waiting_for_new_message: Default::default(),
+            last_waker_waiting_for_new_message: None,
         }
     }
 
     pub fn add_message_to_queue(&mut self, data: Bytes) {
         self.pending_messages.push_back(data);
-        for waker in self.wakers_waiting_for_new_message.drain(..) {
+        if let Some(waker) = self.last_waker_waiting_for_new_message.take() {
             waker.wake();
         }
     }
@@ -76,7 +76,7 @@ impl InboundSession {
             });
             Poll::Ready(())
         } else {
-            self.wakers_waiting_for_new_message.push(cx.waker().clone());
+            self.last_waker_waiting_for_new_message = Some(cx.waker().clone());
             Poll::Pending
         }
     }

crates/apollo_network/src/utils.rs

@@ -13,14 +13,14 @@ use libp2p::{Multiaddr, PeerId};
 // Stream trait from tokio_stream and not from futures.
 pub struct StreamMap<K: Unpin + Clone + Ord, V: Stream + Unpin> {
     map: BTreeMap<K, V>,
-    wakers_waiting_for_new_stream: Vec<Waker>,
+    last_waker_waiting_for_new_stream: Option<Waker>,
     next_index_to_poll: Option<usize>,
 }
 
 impl<K: Unpin + Clone + Ord, V: Stream + Unpin> StreamMap<K, V> {
     #[allow(dead_code)]
     pub fn new(map: BTreeMap<K, V>) -> Self {
-        Self { map, wakers_waiting_for_new_stream: Default::default(), next_index_to_poll: None }
+        Self { map, last_waker_waiting_for_new_stream: None, next_index_to_poll: None }
     }
 
     #[allow(dead_code)]
@@ -40,7 +40,7 @@ impl<K: Unpin + Clone + Ord, V: Stream + Unpin> StreamMap<K, V> {
 
     pub fn insert(&mut self, key: K, value: V) -> Option<V> {
         let res = self.map.insert(key, value);
-        for waker in self.wakers_waiting_for_new_stream.drain(..) {
+        if let Some(waker) = self.last_waker_waiting_for_new_stream.take() {
             waker.wake();
         }
         res
@@ -89,7 +89,7 @@ impl<K: Unpin + Clone + Ord, V: Stream + Unpin> Stream for StreamMap<K, V> {
             unpinned_self.map.remove(&finished_stream_key);
             return Poll::Ready(Some((finished_stream_key, None)));
         }
-        unpinned_self.wakers_waiting_for_new_stream.push(cx.waker().clone());
+        unpinned_self.last_waker_waiting_for_new_stream = Some(cx.waker().clone());
         Poll::Pending
     }
 }