refactor(iroh): Spawn RemoteStateActor in JoinSet and remove guarded_channel.rs

iroh/src/magicsock/remote_map.rs

@@ -2,20 +2,22 @@ use std::{
     collections::{BTreeSet, hash_map},
     hash::Hash,
     net::{IpAddr, SocketAddr},
+    ops::DerefMut,
     sync::{Arc, Mutex},
     time::Duration,
 };
 
 use iroh_base::{EndpointId, RelayUrl};
+use n0_future::task::JoinSet;
 use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
-use tracing::warn;
+use tracing::{debug, error, warn};
 
 pub(crate) use self::remote_state::PathsWatcher;
+use self::remote_state::RemoteStateActor;
 pub(super) use self::remote_state::RemoteStateMessage;
 pub use self::remote_state::{PathInfo, PathInfoList};
-use self::remote_state::{RemoteStateActor, RemoteStateHandle};
 use super::{
     DirectAddr, DiscoState, MagicsockMetrics,
     mapped_addrs::{AddrMap, EndpointIdMappedAddr, RelayMappedAddr},
@@ -45,7 +47,7 @@ pub(crate) struct RemoteMap {
     // State we keep about remote endpoints.
     //
     /// The actors tracking each remote endpoint.
-    actor_handles: Mutex<FxHashMap<EndpointId, RemoteStateHandle>>,
+    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.
@@ -57,10 +59,12 @@ pub(crate) struct RemoteMap {
     /// The endpoint ID of the local endpoint.
     local_endpoint_id: EndpointId,
     metrics: Arc<MagicsockMetrics>,
-    local_addrs: n0_watcher::Direct<BTreeSet<DirectAddr>>,
+    /// The "direct" addresses known for our local endpoint
+    local_direct_addrs: n0_watcher::Direct<BTreeSet<DirectAddr>>,
     disco: DiscoState,
     sender: TransportsSender,
     discovery: ConcurrentDiscovery,
+    actor_tasks: Mutex<JoinSet<(EndpointId, Vec<RemoteStateMessage>)>>,
 }
 
 impl RemoteMap {
@@ -69,21 +73,22 @@ impl RemoteMap {
         local_endpoint_id: EndpointId,
         metrics: Arc<MagicsockMetrics>,
 
-        local_addrs: n0_watcher::Direct<BTreeSet<DirectAddr>>,
+        local_direct_addrs: n0_watcher::Direct<BTreeSet<DirectAddr>>,
         disco: DiscoState,
         sender: TransportsSender,
         discovery: ConcurrentDiscovery,
     ) -> Self {
         Self {
-            actor_handles: Mutex::new(FxHashMap::default()),
+            actor_senders: Mutex::new(FxHashMap::default()),
             endpoint_mapped_addrs: Default::default(),
             relay_mapped_addrs: Default::default(),
             local_endpoint_id,
             metrics,
-            local_addrs,
+            local_direct_addrs,
             disco,
             sender,
             discovery,
+            actor_tasks: Default::default(),
         }
     }
 
@@ -96,8 +101,34 @@ impl RemoteMap {
     /// This should be called periodically to remove handles to endpoint state actors
     /// that have shutdown after their idle timeout expired.
     pub(super) fn remove_closed_remote_state_actors(&self) {
-        let mut handles = self.actor_handles.lock().expect("poisoned");
-        handles.retain(|_eid, handle| !handle.sender.is_closed())
+        let mut senders = self.actor_senders.lock().expect("poisoned");
+        senders.retain(|_eid, sender| !sender.is_closed());
+        while let Some(result) = self.actor_tasks.lock().expect("poisoned").try_join_next() {
+            match result {
+                Ok((eid, leftover_msgs)) => {
+                    let entry = senders.entry(eid);
+                    if leftover_msgs.is_empty() {
+                        match entry {
+                            hash_map::Entry::Occupied(occupied_entry) => occupied_entry.remove(),
+                            hash_map::Entry::Vacant(_) => {
+                                panic!("this should be impossible TODO(matheus23)");
+                            }
+                        };
+                    } else {
+                        // The remote actor got messages while it was closing, so we're restarting
+                        debug!(%eid, "restarting terminated remote state actor: messages received during shutdown");
+                        let sender = self.start_remote_state_actor(eid, leftover_msgs);
+                        entry.insert_entry(sender);
+                    }
+                }
+                Err(err) => {
+                    if let Ok(panic) = err.try_into_panic() {
+                        error!("RemoteStateActor panicked.");
+                        std::panic::resume_unwind(panic);
+                    }
+                }
+            }
+        }
     }
 
     /// Returns the sender for the [`RemoteStateActor`].
@@ -106,21 +137,12 @@ impl RemoteMap {
     ///
     /// [`RemoteStateActor`]: remote_state::RemoteStateActor
     pub(super) fn remote_state_actor(&self, eid: EndpointId) -> mpsc::Sender<RemoteStateMessage> {
-        let mut handles = self.actor_handles.lock().expect("poisoned");
+        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().sender.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, vec![]);
+                entry.insert(sender.clone());
                 sender
             }
         }
@@ -132,22 +154,24 @@ impl RemoteMap {
     fn start_remote_state_actor(
         &self,
         eid: EndpointId,
-    ) -> (RemoteStateHandle, mpsc::Sender<RemoteStateMessage>) {
+        initial_msgs: Vec<RemoteStateMessage>,
+    ) -> mpsc::Sender<RemoteStateMessage> {
         // Ensure there is a RemoteMappedAddr for this EndpointId.
         self.endpoint_mapped_addrs.get(&eid);
-        let handle = RemoteStateActor::new(
+        RemoteStateActor::new(
             eid,
             self.local_endpoint_id,
-            self.local_addrs.clone(),
+            self.local_direct_addrs.clone(),
             self.disco.clone(),
             self.relay_mapped_addrs.clone(),
             self.metrics.clone(),
             self.sender.clone(),
             self.discovery.clone(),
         )
-        .start();
-        let sender = handle.sender.get().expect("just created");
-        (handle, sender)
+        .start(
+            initial_msgs,
+            self.actor_tasks.lock().expect("poisoned").deref_mut(),
+        )
     }
 
     pub(super) fn handle_ping(&self, msg: disco::Ping, sender: EndpointId, src: transports::Addr) {