refactor: remove the TransportsSenderActor

dignifiedquire · dignifiedquire · commit b4f3dc389ef7 · 2025-11-11T13:59:02.000+01:00
diff --git a/iroh/src/magicsock.rs b/iroh/src/magicsock.rs
@@ -1011,15 +1011,14 @@ impl Handle {
         let (disco, disco_receiver) = DiscoState::new(&secret_key);
 
         let endpoint_map = {
-            let sender = transports.create_sender();
             EndpointMap::new(
                 secret_key.public(),
                 // #[cfg(any(test, feature = "test-utils"))]
                 // path_selection,
                 metrics.magicsock.clone(),
-                sender,
                 direct_addrs.addrs.watch(),
                 disco.clone(),
+                transports.create_sender(),
             )
         };
 
diff --git a/iroh/src/magicsock/endpoint_map.rs b/iroh/src/magicsock/endpoint_map.rs
@@ -6,16 +6,15 @@ use std::{
 };
 
 use iroh_base::{EndpointAddr, EndpointId, RelayUrl};
-use n0_future::task::{self, AbortOnDropHandle};
 use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
-use tracing::{Instrument, info_span, trace, warn};
+use tracing::warn;
 
 use super::{
     DirectAddr, DiscoState, MagicsockMetrics,
     mapped_addrs::{AddrMap, EndpointIdMappedAddr, RelayMappedAddr},
-    transports::{self, OwnedTransmit, TransportsSender},
+    transports::{self, TransportsSender},
 };
 use crate::disco::{self};
 // #[cfg(any(test, feature = "test-utils"))]
@@ -58,10 +57,9 @@ pub(crate) struct EndpointMap {
     /// The endpoint ID of the local endpoint.
     local_endpoint_id: EndpointId,
     metrics: Arc<MagicsockMetrics>,
-    /// Handle to an actor that can send over the transports.
-    transports_handle: TransportsSenderHandle,
     local_addrs: n0_watcher::Direct<BTreeSet<DirectAddr>>,
     disco: DiscoState,
+    sender: TransportsSender,
 }
 
 impl EndpointMap {
@@ -71,19 +69,20 @@ impl EndpointMap {
         // TODO:
         // #[cfg(any(test, feature = "test-utils"))] path_selection: PathSelection,
         metrics: Arc<MagicsockMetrics>,
-        sender: TransportsSender,
+
         local_addrs: n0_watcher::Direct<BTreeSet<DirectAddr>>,
         disco: DiscoState,
+        sender: TransportsSender,
     ) -> Self {
         Self {
             actor_handles: Mutex::new(FxHashMap::default()),
             endpoint_mapped_addrs: Default::default(),
             relay_mapped_addrs: Default::default(),
             local_endpoint_id,
             metrics,
-            transports_handle: TransportsSenderActor::new(sender).start(),
             local_addrs,
             disco,
+            sender,
         }
     }
 
@@ -135,18 +134,17 @@ impl EndpointMap {
             Some(handle) => handle.sender.clone(),
             None => {
                 // Create a new EndpointStateActor and insert it into the endpoint map.
-                let sender = self.transports_handle.inbox.clone();
                 let local_addrs = self.local_addrs.clone();
                 let disco = self.disco.clone();
                 let metrics = self.metrics.clone();
                 let actor = EndpointStateActor::new(
                     eid,
                     self.local_endpoint_id,
-                    sender,
                     local_addrs,
                     disco,
                     self.relay_mapped_addrs.clone(),
                     metrics,
+                    self.sender.clone(),
                 );
                 let handle = actor.start();
                 let sender = handle.sender.clone();
@@ -282,79 +280,6 @@ impl IpPort {
     }
 }
 
-/// An actor that can send datagrams onto iroh transports.
-///
-/// The [`EndpointStateActor`]s want to be able to send datagrams.  Because we can not create
-/// [`TransportsSender`]s on demand we must share one for the entire [`EndpointMap`], which
-/// lives in this actor.
-///
-/// [`EndpointStateActor`]: endpoint_state::EndpointStateActor
-#[derive(Debug)]
-struct TransportsSenderActor {
-    sender: TransportsSender,
-}
-
-impl TransportsSenderActor {
-    fn new(sender: TransportsSender) -> Self {
-        Self { sender }
-    }
-
-    fn start(self) -> TransportsSenderHandle {
-        // This actor gets an inbox size of exactly 1.  This is the same as if they had the
-        // underlying sender directly: either you can send or not, or you await until you
-        // can.  No need to introduce extra buffering.
-        let (tx, rx) = mpsc::channel(1);
-
-        let task = task::spawn(
-            async move {
-                self.run(rx).await;
-            }
-            .instrument(info_span!("TransportsSenderActor")),
-        );
-        TransportsSenderHandle {
-            inbox: tx,
-            _task: AbortOnDropHandle::new(task),
-        }
-    }
-
-    async fn run(self, mut inbox: mpsc::Receiver<TransportsSenderMessage>) {
-        use TransportsSenderMessage::SendDatagram;
-
-        while let Some(SendDatagram(dst, owned_transmit)) = inbox.recv().await {
-            let transmit = transports::Transmit {
-                ecn: owned_transmit.ecn,
-                contents: owned_transmit.contents.as_ref(),
-                segment_size: owned_transmit.segment_size,
-            };
-            let len = transmit.contents.len();
-            match self.sender.send(&dst, None, &transmit).await {
-                Ok(()) => {}
-                Err(err) => {
-                    trace!(?dst, %len, "transmit failed to send: {err:#}");
-                }
-            };
-        }
-        trace!("actor terminating");
-    }
-}
-
-#[derive(Debug)]
-struct TransportsSenderHandle {
-    inbox: mpsc::Sender<TransportsSenderMessage>,
-    _task: AbortOnDropHandle<()>,
-}
-
-#[derive(Debug)]
-enum TransportsSenderMessage {
-    SendDatagram(transports::Addr, OwnedTransmit),
-}
-
-impl From<(transports::Addr, OwnedTransmit)> for TransportsSenderMessage {
-    fn from(source: (transports::Addr, OwnedTransmit)) -> Self {
-        Self::SendDatagram(source.0, source.1)
-    }
-}
-
 #[cfg(test)]
 mod tests {
 
diff --git a/iroh/src/magicsock/endpoint_map/endpoint_state.rs b/iroh/src/magicsock/endpoint_map/endpoint_state.rs
@@ -7,7 +7,6 @@ use std::{
 };
 
 use iroh_base::{EndpointAddr, EndpointId, RelayUrl, TransportAddr};
-use n0_error::StdResultExt;
 use n0_future::{
     FuturesUnordered, MergeUnbounded, Stream, StreamExt,
     task::{self, AbortOnDropHandle},
@@ -23,7 +22,7 @@ use tokio::sync::{mpsc, oneshot};
 use tokio_stream::wrappers::{BroadcastStream, errors::BroadcastStreamRecvError};
 use tracing::{Instrument, Level, debug, error, event, info_span, instrument, trace, warn};
 
-use super::{Source, TransportsSenderMessage, path_state::PathState};
+use super::{Source, path_state::PathState};
 // TODO: Use this
 // #[cfg(any(test, feature = "test-utils"))]
 // use crate::endpoint::PathSelection;
@@ -33,7 +32,7 @@ use crate::{
     magicsock::{
         DiscoState, HEARTBEAT_INTERVAL, MagicsockMetrics, PATH_MAX_IDLE_TIMEOUT,
         mapped_addrs::{AddrMap, MappedAddr, RelayMappedAddr},
-        transports::{self, OwnedTransmit},
+        transports::{self, OwnedTransmit, TransportsSender},
     },
     util::MaybeFuture,
 };
@@ -95,10 +94,7 @@ pub(super) struct EndpointStateActor {
     //
     /// Metrics.
     metrics: Arc<MagicsockMetrics>,
-    /// Allowing us to directly send datagrams.
-    ///
-    /// Used for handling [`EndpointStateMessage::SendDatagram`] messages.
-    transports_sender: mpsc::Sender<TransportsSenderMessage>,
+    sender: TransportsSender,
     /// Our local addresses.
     ///
     /// These are our local addresses and any reflexive transport addresses.
@@ -149,17 +145,16 @@ impl EndpointStateActor {
     pub(super) fn new(
         endpoint_id: EndpointId,
         local_endpoint_id: EndpointId,
-        transports_sender: mpsc::Sender<TransportsSenderMessage>,
         local_addrs: n0_watcher::Direct<BTreeSet<DirectAddr>>,
         disco: DiscoState,
         relay_mapped_addrs: AddrMap<(RelayUrl, EndpointId), RelayMappedAddr>,
         metrics: Arc<MagicsockMetrics>,
+        sender: TransportsSender,
     ) -> Self {
         Self {
             endpoint_id,
             local_endpoint_id,
             metrics,
-            transports_sender,
             local_addrs,
             relay_mapped_addrs,
             disco,
@@ -172,6 +167,7 @@ impl EndpointStateActor {
             scheduled_holepunch: None,
             scheduled_open_path: None,
             pending_open_paths: VecDeque::new(),
+            sender,
         }
     }
 
@@ -297,26 +293,34 @@ impl EndpointStateActor {
         Ok(())
     }
 
+    async fn send_datagram(
+        &self,
+        dst: transports::Addr,
+        owned_transmit: OwnedTransmit,
+    ) -> n0_error::Result<()> {
+        let transmit = transports::Transmit {
+            ecn: owned_transmit.ecn,
+            contents: owned_transmit.contents.as_ref(),
+            segment_size: owned_transmit.segment_size,
+        };
+        self.sender.send(&dst, None, &transmit).await?;
+        Ok(())
+    }
+
     /// Handles [`EndpointStateMessage::SendDatagram`].
     ///
     /// Error returns are fatal and kill the actor.
     async fn handle_msg_send_datagram(&mut self, transmit: OwnedTransmit) -> n0_error::Result<()> {
         if let Some(addr) = self.selected_path.get() {
             trace!(?addr, "sending datagram to selected path");
-            self.transports_sender
-                .send((addr, transmit).into())
-                .await
-                .std_context("TransportSenderActor stopped")?;
+            self.send_datagram(addr, transmit).await?;
         } else {
             trace!(
                 paths = ?self.paths.keys().collect::<Vec<_>>(),
                 "sending datagram to all known paths",
             );
             for addr in self.paths.keys() {
-                self.transports_sender
-                    .send((addr.clone(), transmit.clone()).into())
-                    .await
-                    .std_context("TransportSenerActor stopped")?;
+                self.send_datagram(addr.clone(), transmit.clone()).await?;
             }
             // This message is received *before* a connection is added.  So we do
             // not yet have a connection to holepunch.  Instead we trigger
@@ -711,7 +715,7 @@ impl EndpointStateActor {
             transports::Addr::Ip(_) => &self.metrics.send_disco_udp,
             transports::Addr::Relay(_, _) => &self.metrics.send_disco_relay,
         };
-        match self.transports_sender.send((dst, transmit).into()).await {
+        match self.send_datagram(dst, transmit).await {
             Ok(()) => {
                 trace!("sent");
                 counter.inc();
diff --git a/iroh/src/magicsock/transports.rs b/iroh/src/magicsock/transports.rs
@@ -412,7 +412,7 @@ impl Addr {
 }
 
 /// A sender that sends to all our transports.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct TransportsSender {
     #[cfg(not(wasm_browser))]
     ip: Vec<IpSender>,
@@ -602,8 +602,8 @@ impl quinn::AsyncUdpSocket for MagicTransport {
 /// This is special in that it handles [`MultipathMappedAddr::Mixed`] by delegating to the
 /// [`MagicSock`] which expands it back to one or more [`Addr`]s and sends it
 /// using the underlying [`Transports`].
-// TODO: Can I just send the TransportsSender along in the NodeStateMessage::SendDatagram
-// message??  That way you don't have to hook up the sender into the NodeMap!
+// TODO: Can I just send the TransportsSender along in the EndpointStateMessage::SendDatagram
+// message??  That way you don't have to hook up the sender into the EndpointMap!
 #[derive(Debug)]
 #[pin_project::pin_project]
 pub(crate) struct MagicSender {
diff --git a/iroh/src/magicsock/transports/ip.rs b/iroh/src/magicsock/transports/ip.rs
@@ -101,10 +101,12 @@ impl IpTransport {
     }
 
     pub(super) fn create_sender(&self) -> IpSender {
-        let sender = self.socket.clone().create_sender();
+        let socket = self.socket.clone();
+        let sender = socket.clone().create_sender();
         IpSender {
             bind_addr: self.bind_addr,
             sender,
+            socket,
             metrics: self.metrics.clone(),
         }
     }
@@ -136,11 +138,23 @@ impl IpNetworkChangeSender {
 #[pin_project]
 pub(super) struct IpSender {
     bind_addr: SocketAddr,
+    socket: Arc<UdpSocket>,
     #[pin]
     sender: UdpSender,
     metrics: Arc<MagicsockMetrics>,
 }
 
+impl Clone for IpSender {
+    fn clone(&self) -> Self {
+        Self {
+            bind_addr: self.bind_addr,
+            socket: self.socket.clone(),
+            sender: self.socket.clone().create_sender(),
+            metrics: self.metrics.clone(),
+        }
+    }
+}
+
 impl IpSender {
     pub(super) fn is_valid_send_addr(&self, dst: &SocketAddr) -> bool {
         // Our net-tools crate binds sockets to their specific family.  This means an IPv6
