fix(comms): minor edge-case fix to handle inbound connection while di…

…aling
tari-project · Feb 1, 2022 · c0bb3a3 · c0bb3a3
1 parent d84e0fa
commit c0bb3a3
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 7 deletions.
diff --git a/comms/src/connection_manager/dialer.rs b/comms/src/connection_manager/dialer.rs
@@ -73,6 +73,7 @@ pub(crate) enum DialerRequest {
         Option<oneshot::Sender<Result<PeerConnection, ConnectionManagerError>>>,
     ),
     CancelPendingDial(NodeId),
+    NotifyNewInboundConnection(PeerConnection),
 }
 
 pub struct Dialer<TTransport, TBackoff> {
@@ -168,13 +169,29 @@ where
                 self.handle_dial_peer_request(pending_dials, peer, reply_tx);
             },
             CancelPendingDial(peer_id) => {
-                if let Some(mut s) = self.cancel_signals.remove(&peer_id) {
-                    let _ = s.trigger();
+                self.cancel_dial(&peer_id);
+            },
+
+            NotifyNewInboundConnection(conn) => {
+                if conn.is_connected() {
+                    self.resolve_pending_dials(conn);
                 }
             },
         }
     }
 
+    fn cancel_dial(&mut self, peer_id: &NodeId) {
+        if let Some(mut s) = self.cancel_signals.remove(&peer_id) {
+            let _ = s.trigger();
+        }
+    }
+
+    fn resolve_pending_dials(&mut self, conn: PeerConnection) {
+        let peer = conn.peer_node_id().clone();
+        self.reply_to_pending_requests(&peer, Ok(conn));
+        self.cancel_dial(&peer);
+    }
+
     fn is_pending_dial(&self, node_id: &NodeId) -> bool {
         self.cancel_signals.contains_key(node_id)
     }
@@ -223,12 +240,9 @@ where
             );
         }
 
-        if self.pending_dial_requests.contains_key(&node_id) {
-            self.reply_to_pending_requests(&node_id, dial_result);
-        }
+        self.reply_to_pending_requests(&node_id, dial_result);
 
-        // Drop cancel signal
-        let _ = self.cancel_signals.remove(&node_id);
+        self.cancel_dial(&node_id);
     }
 
     pub async fn notify_connection_manager(&mut self, event: ConnectionManagerEvent) {

diff --git a/comms/src/connection_manager/manager.rs b/comms/src/connection_manager/manager.rs
@@ -424,6 +424,13 @@ where
             },
 
             PeerConnected(conn) => {
+                if conn.direction().is_inbound() {
+                    // Notify the dialer that we have an inbound connection, so that is can resolve any pending dials.
+                    let _ = self
+                        .dialer_tx
+                        .send(DialerRequest::NotifyNewInboundConnection(conn.clone()))
+                        .await;
+                }
                 metrics::successful_connections(conn.peer_node_id(), conn.direction()).inc();
                 self.publish_event(PeerConnected(conn));
             },