Merge pull request #3300 from nymtech/bugfix/empty-ack-stream-map

make sure to clear inner 'ack_map' in 'GatewaysReader'

Author: tommyv1987

nym-api/src/network_monitor/gateways_reader.rs

@@ -1,8 +1,11 @@
 // Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
+use futures::Stream;
 use nym_crypto::asymmetric::identity;
 use nym_gateway_client::{AcknowledgementReceiver, MixnetMessageReceiver};
+use std::pin::Pin;
+use std::task::{Context, Poll};
 use tokio_stream::StreamMap;
 
 pub(crate) type GatewayMessages = Vec<Vec<u8>>;
@@ -20,11 +23,7 @@ impl GatewaysReader {
         }
     }
 
-    pub fn stream_map(&mut self) -> &mut StreamMap<String, MixnetMessageReceiver> {
-        &mut self.stream_map
-    }
-
-    pub fn add_recievers(
+    pub fn add_receivers(
         &mut self,
         id: identity::PublicKey,
         message_receiver: MixnetMessageReceiver,
@@ -35,8 +34,27 @@ impl GatewaysReader {
         self.ack_map.insert(channel_id, ack_receiver);
     }
 
-    pub fn remove_recievers(&mut self, id: &str) {
+    pub fn remove_receivers(&mut self, id: &str) {
         self.stream_map.remove(id);
         self.ack_map.remove(id);
     }
 }
+
+impl Stream for GatewaysReader {
+    // just return whatever is returned by our main `stream_map`
+    type Item = <StreamMap<String, MixnetMessageReceiver> as Stream>::Item;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        // exhaust the ack map if possible
+        match Pin::new(&mut self.ack_map).poll_next(cx) {
+            Poll::Ready(None) => {
+                // this should have never happened!
+                return Poll::Ready(None);
+            }
+            Poll::Ready(Some(_item)) => (),
+            Poll::Pending => (),
+        }
+
+        Pin::new(&mut self.stream_map).poll_next(cx)
+    }
+}

nym-api/src/network_monitor/monitor/receiver.rs

@@ -42,10 +42,10 @@ impl PacketReceiver {
         match update {
             GatewayClientUpdate::New(id, (message_receiver, ack_receiver)) => {
                 self.gateways_reader
-                    .add_recievers(id, message_receiver, ack_receiver);
+                    .add_receivers(id, message_receiver, ack_receiver);
             }
             GatewayClientUpdate::Failure(id) => {
-                self.gateways_reader.remove_recievers(&id.to_string());
+                self.gateways_reader.remove_receivers(&id.to_string());
             }
         }
     }
@@ -66,11 +66,12 @@ impl PacketReceiver {
                 // unwrap here is fine as it can only return a `None` if the PacketSender has died
                 // and if that was the case, then the entire monitor is already in an undefined state
                 update = self.clients_updater.next() => self.process_gateway_update(update.unwrap()),
-                // similarly gateway reader will never return a `None` as it's implemented
-                // as an infinite stream that returns Poll::Pending if it doesn't have anything
-                // to return
-                Some((_gateway_id, message)) = self.gateways_reader.stream_map().next() => {
-                        self.process_gateway_messages(message)
+                gateway_message = self.gateways_reader.next() => {
+                    let Some((_gateway_id, message)) = gateway_message else {
+                        log::error!("the gateways reader stream has terminated!");
+                        continue
+                    };
+                    self.process_gateway_messages(message)
                 }
             }
         }