cleanup

Author: Frando

src/net.rs

@@ -287,13 +287,14 @@ type AcceptRemoteRequestsStream =
 
 struct Actor {
     me: NodeId,
+    endpoint: Endpoint,
     alpn: Bytes,
     config: Config,
     local_rx: mpsc::Receiver<LocalActorMessage>,
     local_tx: mpsc::Sender<LocalActorMessage>,
     api_rx: mpsc::Receiver<api::RpcMessage>,
     topics: HashMap<TopicId, TopicHandle>,
-    topic_pending_remotes: HashMap<NodeId, HashSet<TopicId>>,
+    pending_remotes_with_topics: HashMap<NodeId, HashSet<TopicId>>,
     topic_tasks: JoinSet<TopicActor>,
     remotes: HashMap<NodeId, RemoteState>,
     close_connections: JoinSet<(NodeId, Connection)>,
@@ -332,18 +333,19 @@ impl Actor {
             api_tx,
             local_tx.clone(),
             Actor {
+                endpoint,
                 me,
                 config,
                 api_rx,
                 local_tx,
                 local_rx,
                 node_addr_updates: Box::pin(node_addr_updates),
-                dialer: Dialer::new(endpoint),
+                dialer: Dialer::default(),
                 our_peer_data: Watchable::new(initial_peer_data),
                 alpn: alpn.unwrap_or_else(|| crate::ALPN.to_vec().into()),
                 metrics: metrics.clone(),
                 topics: Default::default(),
-                topic_pending_remotes: Default::default(),
+                pending_remotes_with_topics: Default::default(),
                 remotes: Default::default(),
                 close_connections: JoinSet::new(),
                 topic_tasks: JoinSet::new(),
@@ -414,7 +416,7 @@ impl Actor {
                 }
                 true
             }
-            (node_id, res) = self.dialer.next_conn() => {
+            Some((node_id, res)) = self.dialer.next(), if !self.dialer.is_empty() => {
                 trace!(remote=%node_id.fmt_short(), ok=res.is_ok(), "tick: dialed");
                 self.handle_remote_connection(node_id, res, Direction::Dial).await;
                 true
@@ -464,14 +466,14 @@ impl Actor {
 
     #[cfg(test)]
     fn endpoint(&self) -> &Endpoint {
-        self.dialer.endpoint()
+        &self.endpoint
     }
 
     fn drain_pending_dials(
         &mut self,
         remote: &NodeId,
     ) -> impl Iterator<Item = (TopicId, &TopicHandle)> {
-        self.topic_pending_remotes
+        self.pending_remotes_with_topics
             .remove(remote)
             .into_iter()
             .flatten()
@@ -485,13 +487,15 @@ impl Actor {
         if let Some(state) = self.remotes.get(&remote) {
             let tx = handle.tx.clone();
             let state = state.clone();
+            // TODO: Track task?
             task::spawn(async move {
                 let msg = state.open_topic(topic_id).await;
                 tx.send(msg).await.ok();
             });
         } else {
-            self.dialer.queue_dial(remote, self.alpn.clone());
-            self.topic_pending_remotes
+            self.dialer
+                .queue_dial(&self.endpoint, remote, self.alpn.clone());
+            self.pending_remotes_with_topics
                 .entry(remote)
                 .or_default()
                 .insert(topic_id);

src/net/dialer.rs

@@ -1,80 +1,46 @@
-use std::collections::HashMap;
+use std::collections::HashSet;
 
 use bytes::Bytes;
 use iroh::{
     endpoint::{ConnectError, Connection},
     Endpoint, NodeId,
 };
 use tokio::task::JoinSet;
-use tracing::{error, Instrument};
+use tracing::Instrument;
 
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub(crate) struct Dialer {
-    endpoint: Endpoint,
-    pending: JoinSet<(NodeId, Result<Connection, ConnectError>)>,
-    pending_dials: HashMap<NodeId, ()>,
+    pending_tasks: JoinSet<(NodeId, Result<Connection, ConnectError>)>,
+    pending_nodes: HashSet<NodeId>,
 }
 
 impl Dialer {
-    /// Create a new dialer for a [`Endpoint`]
-    pub(crate) fn new(endpoint: Endpoint) -> Self {
-        Self {
-            endpoint,
-            pending: Default::default(),
-            pending_dials: Default::default(),
-        }
-    }
-
-    #[cfg(test)]
-    pub(crate) fn endpoint(&self) -> &Endpoint {
-        &self.endpoint
-    }
-
     /// Starts to dial a node by [`NodeId`].
-    pub(crate) fn queue_dial(&mut self, node_id: NodeId, alpn: Bytes) {
-        if self.is_pending(node_id) {
-            return;
+    pub(crate) fn queue_dial(&mut self, endpoint: &Endpoint, node_id: NodeId, alpn: Bytes) {
+        if self.pending_nodes.insert(node_id) {
+            let endpoint = endpoint.clone();
+            let fut = async move { (node_id, endpoint.connect(node_id, &alpn).await) }
+                .instrument(tracing::Span::current());
+            self.pending_tasks.spawn(fut);
         }
-        self.pending_dials.insert(node_id, ());
-        let endpoint = self.endpoint.clone();
-        self.pending.spawn(
-            async move {
-                let res = endpoint.connect(node_id, &alpn).await;
-                (node_id, res)
-            }
-            .instrument(tracing::Span::current()),
-        );
     }
 
-    /// Checks if a node is currently being dialed.
-    pub(crate) fn is_pending(&self, node: NodeId) -> bool {
-        self.pending_dials.contains_key(&node)
+    pub(crate) fn is_empty(&self) -> bool {
+        self.pending_tasks.is_empty()
     }
 
     /// Waits for the next dial operation to complete.
     /// `None` means disconnected
-    pub(crate) async fn next_conn(&mut self) -> (NodeId, Result<Connection, ConnectError>) {
-        match self.pending_dials.is_empty() {
-            false => {
-                let (node_id, res) = loop {
-                    match self.pending.join_next().await {
-                        Some(Ok((node_id, res))) => {
-                            self.pending_dials.remove(&node_id);
-                            break (node_id, res);
-                        }
-                        Some(Err(e)) => {
-                            error!("next conn error: {:?}", e);
-                        }
-                        None => {
-                            error!("no more pending conns available");
-                            std::future::pending().await
-                        }
-                    }
-                };
-
-                (node_id, res)
+    ///
+    /// Will be pending forever if no connections are in progress.
+    pub(crate) async fn next(&mut self) -> Option<(NodeId, Result<Connection, ConnectError>)> {
+        match self.pending_tasks.join_next().await {
+            Some(res) => {
+                let (node_id, res) = res.expect("connect task panicked");
+                self.pending_nodes.remove(&node_id);
+                Some((node_id, res))
             }
-            true => std::future::pending().await,
+            None => None,
         }
     }
 }