Replace RemoteStateHandle with storing a JoinSet in RemoteMap

Author: matheus23

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 iroh_base::{EndpointAddr, 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, GuardedSender<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.
@@ -61,6 +63,7 @@ pub(crate) struct RemoteMap {
     disco: DiscoState,
     sender: TransportsSender,
     discovery: ConcurrentDiscovery,
+    actor_tasks: Mutex<JoinSet<Vec<RemoteStateMessage>>>,
 }
 
 impl RemoteMap {
@@ -75,7 +78,7 @@ impl RemoteMap {
         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,
@@ -84,6 +87,7 @@ impl RemoteMap {
             disco,
             sender,
             discovery,
+            actor_tasks: Default::default(),
         }
     }
 
@@ -96,8 +100,19 @@ 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(leftover_msgs) => debug!(?leftover_msgs, "TODO: handle leftover messages"),
+                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,10 +121,10 @@ 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() {
+                if let Some(sender) = entry.get().get() {
                     sender
                 } else {
                     // The actor is dead: Start a new actor.
@@ -132,10 +147,13 @@ impl RemoteMap {
     fn start_remote_state_actor(
         &self,
         eid: EndpointId,
-    ) -> (RemoteStateHandle, mpsc::Sender<RemoteStateMessage>) {
+    ) -> (
+        GuardedSender<RemoteStateMessage>,
+        mpsc::Sender<RemoteStateMessage>,
+    ) {
         // Ensure there is a RemoteMappedAddr for this EndpointId.
         self.endpoint_mapped_addrs.get(&eid);
-        let handle = RemoteStateActor::new(
+        let sender = RemoteStateActor::new(
             eid,
             self.local_endpoint_id,
             self.local_addrs.clone(),
@@ -145,9 +163,9 @@ impl RemoteMap {
             self.sender.clone(),
             self.discovery.clone(),
         )
-        .start();
-        let sender = handle.sender.get().expect("just created");
-        (handle, sender)
+        .start(self.actor_tasks.lock().expect("poisoned").deref_mut());
+        let tx = sender.get().expect("just created");
+        (sender, tx)
     }
 
     pub(super) fn handle_ping(&self, msg: disco::Ping, sender: EndpointId, src: transports::Addr) {

iroh/src/magicsock/remote_map/remote_state.rs

@@ -10,7 +10,7 @@ use iroh_base::{EndpointId, RelayUrl, TransportAddr};
 use n0_future::{
     Either, FuturesUnordered, MergeUnbounded, Stream, StreamExt,
     boxed::BoxStream,
-    task::{self, AbortOnDropHandle},
+    task::JoinSet,
     time::{self, Duration, Instant},
 };
 use n0_watcher::{Watchable, Watcher};
@@ -41,7 +41,7 @@ use crate::{
     util::MaybeFuture,
 };
 
-mod guarded_channel;
+pub(crate) mod guarded_channel;
 mod path_state;
 
 // TODO: use this
@@ -209,7 +209,10 @@ impl RemoteStateActor {
         }
     }
 
-    pub(super) fn start(self) -> RemoteStateHandle {
+    pub(super) fn start(
+        self,
+        tasks: &mut JoinSet<Vec<RemoteStateMessage>>,
+    ) -> GuardedSender<RemoteStateMessage> {
         let (tx, rx) = guarded_channel(16);
         let me = self.local_endpoint_id;
         let endpoint_id = self.endpoint_id;
@@ -219,28 +222,28 @@ impl RemoteStateActor {
         // we don't explicitly set a span we get the spans from whatever call happens to
         // first create the actor, which is often very confusing as it then keeps those
         // spans for all logging of the actor.
-        let task = task::spawn(self.run(rx).instrument(info_span!(
+        tasks.spawn(self.run(rx).instrument(info_span!(
             parent: None,
             "RemoteStateActor",
             me = %me.fmt_short(),
             remote = %endpoint_id.fmt_short(),
         )));
-        RemoteStateHandle {
-            sender: tx,
-            _task: AbortOnDropHandle::new(task),
-        }
+        tx
     }
 
     /// Runs the main loop of the actor.
     ///
     /// Note that the actor uses async handlers for tasks from the main loop.  The actor is
     /// not processing items from the inbox while waiting on any async calls.  So some
     /// discipline is needed to not turn pending for a long time.
-    async fn run(mut self, mut inbox: GuardedReceiver<RemoteStateMessage>) {
+    async fn run(
+        mut self,
+        mut inbox: GuardedReceiver<RemoteStateMessage>,
+    ) -> Vec<RemoteStateMessage> {
         trace!("actor started");
         let idle_timeout = time::sleep(ACTOR_MAX_IDLE_TIMEOUT);
         n0_future::pin!(idle_timeout);
-        loop {
+        let leftover_msgs = loop {
             let scheduled_path_open = match self.scheduled_open_path {
                 Some(when) => MaybeFuture::Some(time::sleep_until(when)),
                 None => MaybeFuture::None,
@@ -261,7 +264,7 @@ impl RemoteStateActor {
                 msg = inbox.recv() => {
                     match msg {
                         Some(msg) => self.handle_message(msg).await,
-                        None => break,
+                        None => break vec![],
                     }
                 }
                 Some((id, evt)) = self.path_events.next() => {
@@ -297,15 +300,16 @@ impl RemoteStateActor {
                 _ = &mut idle_timeout => {
                     if self.connections.is_empty() && inbox.close_if_idle() {
                         trace!("idle timeout expired and still idle: terminate actor");
-                        break;
+                        break vec![];
                     } 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
     }
 
     /// Handles an actor message.
@@ -1096,21 +1100,6 @@ pub(crate) enum RemoteStateMessage {
     Latency(oneshot::Sender<Option<Duration>>),
 }
 
-/// A handle to a [`RemoteStateActor`].
-///
-/// Dropping this will stop the actor. The actor will also stop after an idle timeout
-/// if it has no connections, an empty inbox, and no other senders than the one stored
-/// in the endpoint map exist.
-#[derive(Debug)]
-pub(super) struct RemoteStateHandle {
-    /// Sender for the channel into the [`RemoteStateActor`].
-    ///
-    /// This is a [`GuardedSender`], from which we can get a sender but only if the receiver
-    /// hasn't been closed.
-    pub(super) sender: GuardedSender<RemoteStateMessage>,
-    _task: AbortOnDropHandle<()>,
-}
-
 /// Information about a holepunch attempt.
 #[derive(Debug)]
 struct HolepunchAttempt {