Remove guarded_channel and use tokio::mpsc instead. Return leftover messages

Author: matheus23

iroh/src/magicsock/remote_map.rs

@@ -7,7 +7,7 @@ use std::{
     time::Duration,
 };
 
-use iroh_base::{EndpointAddr, EndpointId, RelayUrl};
+use iroh_base::{EndpointId, RelayUrl};
 use n0_future::task::JoinSet;
 use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
@@ -47,7 +47,7 @@ pub(crate) struct RemoteMap {
     // State we keep about remote endpoints.
     //
     /// The actors tracking each remote endpoint.
-    actor_senders: Mutex<FxHashMap<EndpointId, GuardedSender<RemoteStateMessage>>>,
+    actor_senders: Mutex<FxHashMap<EndpointId, mpsc::Sender<RemoteStateMessage>>>,
     /// The mapping between [`EndpointId`]s and [`EndpointIdMappedAddr`]s.
     pub(super) endpoint_mapped_addrs: AddrMap<EndpointId, EndpointIdMappedAddr>,
     /// The mapping between endpoints via a relay and their [`RelayMappedAddr`]s.
@@ -123,19 +123,10 @@ impl RemoteMap {
     pub(super) fn remote_state_actor(&self, eid: EndpointId) -> mpsc::Sender<RemoteStateMessage> {
         let mut handles = self.actor_senders.lock().expect("poisoned");
         match handles.entry(eid) {
-            hash_map::Entry::Occupied(mut entry) => {
-                if let Some(sender) = entry.get().get() {
-                    sender
-                } else {
-                    // The actor is dead: Start a new actor.
-                    let (handle, sender) = self.start_remote_state_actor(eid);
-                    entry.insert(handle);
-                    sender
-                }
-            }
+            hash_map::Entry::Occupied(entry) => entry.get().clone(),
             hash_map::Entry::Vacant(entry) => {
-                let (handle, sender) = self.start_remote_state_actor(eid);
-                entry.insert(handle);
+                let sender = self.start_remote_state_actor(eid);
+                entry.insert(sender.clone());
                 sender
             }
         }
@@ -144,16 +135,10 @@ impl RemoteMap {
     /// Starts a new remote state actor and returns a handle and a sender.
     ///
     /// The handle is not inserted into the endpoint map, this must be done by the caller of this function.
-    fn start_remote_state_actor(
-        &self,
-        eid: EndpointId,
-    ) -> (
-        GuardedSender<RemoteStateMessage>,
-        mpsc::Sender<RemoteStateMessage>,
-    ) {
+    fn start_remote_state_actor(&self, eid: EndpointId) -> mpsc::Sender<RemoteStateMessage> {
         // Ensure there is a RemoteMappedAddr for this EndpointId.
         self.endpoint_mapped_addrs.get(&eid);
-        let sender = RemoteStateActor::new(
+        RemoteStateActor::new(
             eid,
             self.local_endpoint_id,
             self.local_addrs.clone(),
@@ -163,9 +148,7 @@ impl RemoteMap {
             self.sender.clone(),
             self.discovery.clone(),
         )
-        .start(self.actor_tasks.lock().expect("poisoned").deref_mut());
-        let tx = sender.get().expect("just created");
-        (sender, tx)
+        .start(self.actor_tasks.lock().expect("poisoned").deref_mut())
     }
 
     pub(super) fn handle_ping(&self, msg: disco::Ping, sender: EndpointId, src: transports::Addr) {

iroh/src/magicsock/remote_map/remote_state.rs

@@ -19,14 +19,12 @@ use quinn_proto::{PathError, PathEvent, PathId, PathStatus};
 use rustc_hash::FxHashMap;
 use smallvec::SmallVec;
 use sync_wrapper::SyncStream;
-use tokio::sync::oneshot;
 use tokio_stream::wrappers::{BroadcastStream, errors::BroadcastStreamRecvError};
 use tracing::{Instrument, Level, debug, error, event, info_span, instrument, trace, warn};
 
-use self::{
-    guarded_channel::{GuardedReceiver, GuardedSender, guarded_channel},
-    path_state::RemotePathState,
-};
+use self::path_state::RemotePathState;
+use tokio::sync::{mpsc, oneshot};
+
 use super::Source;
 use crate::{
     disco::{self},
@@ -41,7 +39,6 @@ use crate::{
     util::MaybeFuture,
 };
 
-pub(crate) mod guarded_channel;
 mod path_state;
 
 // TODO: use this
@@ -212,8 +209,8 @@ impl RemoteStateActor {
     pub(super) fn start(
         self,
         tasks: &mut JoinSet<Vec<RemoteStateMessage>>,
-    ) -> GuardedSender<RemoteStateMessage> {
-        let (tx, rx) = guarded_channel(16);
+    ) -> mpsc::Sender<RemoteStateMessage> {
+        let (tx, rx) = mpsc::channel(16);
         let me = self.local_endpoint_id;
         let endpoint_id = self.endpoint_id;
 
@@ -238,7 +235,7 @@ impl RemoteStateActor {
     /// discipline is needed to not turn pending for a long time.
     async fn run(
         mut self,
-        mut inbox: GuardedReceiver<RemoteStateMessage>,
+        mut inbox: mpsc::Receiver<RemoteStateMessage>,
     ) -> Vec<RemoteStateMessage> {
         trace!("actor started");
         let idle_timeout = time::sleep(ACTOR_MAX_IDLE_TIMEOUT);
@@ -254,7 +251,7 @@ impl RemoteStateActor {
                 None => MaybeFuture::None,
             };
             n0_future::pin!(scheduled_hp);
-            if !inbox.is_idle() || !self.connections.is_empty() {
+            if !inbox.is_empty() || !self.connections.is_empty() {
                 idle_timeout
                     .as_mut()
                     .reset(Instant::now() + ACTOR_MAX_IDLE_TIMEOUT);
@@ -298,16 +295,22 @@ impl RemoteStateActor {
                     self.handle_discovery_item(item);
                 }
                 _ = &mut idle_timeout => {
-                    if self.connections.is_empty() && inbox.close_if_idle() {
+                    if self.connections.is_empty() && inbox.is_empty() {
                         trace!("idle timeout expired and still idle: terminate actor");
-                        break vec![];
+                        inbox.close();
+                        // There might be a race between checking `inbox.is_empty()` and `inbox.close()`,
+                        // so we pull out all messages that are left over.
+                        let mut leftover_msgs = Vec::with_capacity(inbox.len());
+                        inbox.recv_many(&mut leftover_msgs, inbox.len()).await;
+                        break leftover_msgs;
                     } else {
                         // Seems like we weren't really idle, so we reset
                         idle_timeout.as_mut().reset(Instant::now() + ACTOR_MAX_IDLE_TIMEOUT);
                     }
                 }
             }
         };
+
         trace!("actor terminating");
         leftover_msgs
     }