Better management of disconnecting peers

Rather than send a message when a peer disconnects and then handle that
message asynchronously, we now wait for the task itself to exit and then
return a `DisconnectedPeer` from `ConnectionManager::step`.

We only return the `PeerId` for the `DisconnectedPeer` if it is
still the existing `established connection` for the given peer and
it hasn't been replaced by a newer connection. This prevents calling
`Node::on_disconnect` for the stale connection when it might have
already received an `on_connect` call for the new connection.

Author: andrewjstone

trust-quorum/src/connection_manager.rs

@@ -55,6 +55,19 @@ pub enum MainToConnMsg {
     Msg(WireMsg),
 }
 
+/// The task for this sprockets connection just exited. If
+/// `ConnectionManager::step` returns this value and `peer_id` is `Some` than
+/// it means no new connection for the peer has yet been established. It is
+/// safe to cleanup state for the given `peer_id`, by, for instance, calling
+/// `Node::on_disconnect`.
+///
+/// By always returning the `task_id`, we allow cleanup of proxy requests for
+/// stale nodes that will never complete.
+pub struct DisconnectedPeer {
+    pub task_id: task::Id,
+    pub peer_id: Option<BaseboardId>,
+}
+
 /// All possible messages sent over established connections
 ///
 /// This include trust quorum related `PeerMsg`s, but also ancillary network
@@ -106,7 +119,6 @@ pub enum ConnToMainMsgInner {
     Connected { addr: SocketAddrV6, peer_id: BaseboardId },
     Received { from: BaseboardId, msg: PeerMsg },
     ReceivedNetworkConfig { from: BaseboardId, config: NetworkConfig },
-    Disconnected { peer_id: BaseboardId },
     ProxyRequestReceived { from: BaseboardId, req: proxy::WireRequest },
     ProxyResponseReceived { from: BaseboardId, rsp: proxy::WireResponse },
 }
@@ -154,7 +166,7 @@ impl BiHashItem for TaskHandle {
 }
 
 pub struct EstablishedTaskHandle {
-    baseboard_id: BaseboardId,
+    pub baseboard_id: BaseboardId,
     task_handle: TaskHandle,
 }
 
@@ -246,7 +258,7 @@ pub struct ConnMgrStatus {
 
 /// The state of a proxy connection
 pub enum ProxyConnState {
-    Connected,
+    Connected(task::Id),
     Disconnected,
 }
 
@@ -447,7 +459,7 @@ impl ConnMgr {
         if let Some(h) = self.established.get1(destination) {
             info!(self.log, "Sending {req:?}"; "peer_id" => %destination);
             h.send(req).await;
-            ProxyConnState::Connected
+            ProxyConnState::Connected(h.task_id())
         } else {
             ProxyConnState::Disconnected
         }
@@ -471,32 +483,37 @@ impl ConnMgr {
 
     /// Perform any polling related operations that the connection
     /// manager must perform concurrently.
+    ///
+    /// Return `Ok(Some(DisconnectedPeer))` if an `EstablishedConnectionTask`
+    /// that was still the exclusive connection task for a specific peer has
+    /// just exited.
     pub async fn step(
         &mut self,
         corpus: Vec<Utf8PathBuf>,
-    ) -> Result<(), AcceptError> {
-        tokio::select! {
+    ) -> Result<Option<DisconnectedPeer>, AcceptError> {
+        let disconnected_peer = tokio::select! {
             acceptor = self.server.accept(corpus.clone()) => {
                 self.accept(acceptor?).await?;
+                None
             }
             Some(res) = self.join_set.join_next_with_id() => {
                 match res {
                     Ok((task_id, _)) => {
-                        self.on_task_exit(task_id).await;
+                        Some(self.on_task_exit(task_id))
                     }
                     Err(err) => {
                         warn!(self.log, "Connection task panic: {err}");
-                        self.on_task_exit(err.id()).await;
+                        Some(self.on_task_exit(err.id()))
                     }
-
                 }
             }
             _ = self.reconnect_interval.tick() => {
                 self.reconnect(corpus.clone()).await;
+                None
             }
-        }
+        };
 
-        Ok(())
+        Ok(disconnected_peer)
     }
 
     pub async fn accept(
@@ -686,22 +703,6 @@ impl ConnMgr {
         }
     }
 
-    /// The established connection task has asynchronously exited.
-    pub async fn on_disconnected(
-        &mut self,
-        task_id: task::Id,
-        peer_id: BaseboardId,
-    ) {
-        if let Some(established_task_handle) = self.established.get1(&peer_id) {
-            if established_task_handle.task_id() != task_id {
-                // This was a stale disconnect
-                return;
-            }
-        }
-        warn!(self.log, "peer disconnected"; "peer_id" => %peer_id);
-        let _ = self.established.remove1(&peer_id);
-    }
-
     /// Initiate connections if a corresponding task doesn't already exist. This
     /// must be called periodically to handle transient disconnections which
     /// cause tasks to exit.
@@ -740,9 +741,9 @@ impl ConnMgr {
         &mut self,
         addrs: BTreeSet<SocketAddrV6>,
         corpus: Vec<Utf8PathBuf>,
-    ) -> BTreeSet<BaseboardId> {
+    ) -> Vec<EstablishedTaskHandle> {
         if self.bootstrap_addrs == addrs {
-            return BTreeSet::new();
+            return vec![];
         }
 
         // We don't try to compare addresses from accepted nodes. If DDMD
@@ -770,10 +771,10 @@ impl ConnMgr {
             self.connect_client(corpus.clone(), addr).await;
         }
 
-        let mut disconnected_peers = BTreeSet::new();
+        let mut disconnected_peers = Vec::new();
         for addr in to_disconnect {
-            if let Some(peer_id) = self.disconnect_client(addr).await {
-                disconnected_peers.insert(peer_id);
+            if let Some(handle) = self.disconnect_client(addr).await {
+                disconnected_peers.push(handle);
             }
         }
         disconnected_peers
@@ -861,7 +862,7 @@ impl ConnMgr {
     async fn disconnect_client(
         &mut self,
         addr: SocketAddrV6,
-    ) -> Option<BaseboardId> {
+    ) -> Option<EstablishedTaskHandle> {
         if let Some(handle) = self.connecting.remove2(&addr) {
             // The connection has not yet completed its handshake
             info!(
@@ -880,15 +881,17 @@ impl ConnMgr {
                     "peer_id" => %handle.baseboard_id
                 );
                 handle.abort();
-                Some(handle.baseboard_id)
+                Some(handle)
             } else {
                 None
             }
         }
     }
 
     /// Remove any references to the given task
-    async fn on_task_exit(&mut self, task_id: task::Id) {
+    ///
+    /// Return a `DisconnectedPeer` for the given `task_id`.
+    fn on_task_exit(&mut self, task_id: task::Id) -> DisconnectedPeer {
         // We're most likely to find the task as established so we start with that
         if let Some(handle) = self.established.remove2(&task_id) {
             info!(
@@ -898,6 +901,10 @@ impl ConnMgr {
                 "peer_addr" => %handle.addr(),
                 "peer_id" => %handle.baseboard_id
             );
+            return DisconnectedPeer {
+                task_id,
+                peer_id: Some(handle.baseboard_id),
+            };
         } else if let Some(handle) = self.accepting.remove1(&task_id) {
             info!(
                 self.log,
@@ -919,6 +926,8 @@ impl ConnMgr {
                 "task_id" => ?task_id
             );
         }
+
+        DisconnectedPeer { task_id, peer_id: None }
     }
 }
 

trust-quorum/src/established_conn.rs

@@ -155,14 +155,6 @@ impl EstablishedConn {
     }
 
     async fn close(&mut self) {
-        if let Err(_) = self.main_tx.try_send(ConnToMainMsg {
-            task_id: self.task_id,
-            msg: ConnToMainMsgInner::Disconnected {
-                peer_id: self.peer_id.clone(),
-            },
-        }) {
-            warn!(self.log, "Failed to send to main task");
-        }
         let _ = self.writer.shutdown().await;
     }
 

trust-quorum/src/proxy.rs

@@ -23,6 +23,7 @@ use omicron_uuid_kinds::RackUuid;
 use serde::{Deserialize, Serialize};
 use slog_error_chain::{InlineErrorChain, SlogInlineError};
 use tokio::sync::{mpsc, oneshot};
+use tokio::task;
 use trust_quorum_protocol::{
     BaseboardId, CommitError, Configuration, Epoch, PrepareAndCommitError,
 };
@@ -246,20 +247,27 @@ pub enum TrackerError {
 /// A trackable in-flight proxy request, owned by the `Tracker`
 #[derive(Debug)]
 pub struct TrackableRequest {
+    /// Each `TrackableRequest` is bound to a
+    /// [`crate::established_conn::EstablishedConn`] that is uniquely identified
+    /// by its `tokio::task::Id`. This is useful because it disambiguates
+    /// connect and disconnect operations for the same `destination` such that
+    /// they don't have to be totally ordered. It is enough to know that a
+    /// disconnect for a given `task_id` only occurs after a connect.
+    task_id: task::Id,
+    /// A unique id for a given proxy request
     request_id: Uuid,
-    destination: BaseboardId,
     // The option exists so we can take the sender out in `on_disconnect`, when
     // the request is borrowed, but about to be discarded.
     tx: DebugIgnore<Option<oneshot::Sender<Result<WireValue, TrackerError>>>>,
 }
 
 impl TrackableRequest {
     pub fn new(
-        destination: BaseboardId,
+        task_id: task::Id,
         request_id: Uuid,
         tx: oneshot::Sender<Result<WireValue, TrackerError>>,
     ) -> TrackableRequest {
-        TrackableRequest { request_id, destination, tx: DebugIgnore(Some(tx)) }
+        TrackableRequest { task_id, request_id, tx: DebugIgnore(Some(tx)) }
     }
 }
 
@@ -306,9 +314,9 @@ impl Tracker {
     }
 
     /// A remote peer has disconnected
-    pub fn on_disconnect(&mut self, from: &BaseboardId) {
+    pub fn on_disconnect(&mut self, task_id: task::Id) {
         self.ops.retain(|mut req| {
-            if &req.destination == from {
+            if req.task_id == task_id {
                 let tx = req.tx.take().unwrap();
                 let _ = tx.send(Err(TrackerError::Disconnected));
                 false
@@ -335,15 +343,15 @@ mod tests {
     async fn recv_and_insert(
         rx: &mut mpsc::Receiver<NodeApiRequest>,
         tracker: &mut Tracker,
+        task_id: task::Id,
     ) {
-        let Some(NodeApiRequest::Proxy { destination, wire_request, tx }) =
+        let Some(NodeApiRequest::Proxy { wire_request, tx, .. }) =
             rx.recv().await
         else {
             panic!("Invalid NodeApiRequest")
         };
 
-        let req =
-            TrackableRequest::new(destination, wire_request.request_id, tx);
+        let req = TrackableRequest::new(task_id, wire_request.request_id, tx);
         tracker.insert(req);
     }
 
@@ -355,6 +363,11 @@ mod tests {
         };
         let rack_id = RackUuid::new_v4();
 
+        // In real code, the `tokio::task::ID` is the id of the
+        // `EstablishedConnectionTask`. However, we are simulating those
+        // connections here, so just use an ID of an arbitrary task.
+        let task_id = task::spawn(async {}).id();
+
         // Test channel where the sender is usually cloned from the [`crate::NodeTaskHandle`],
         // and the receiver is owned by the local [`crate::NodeTask`]
         let (tx, mut rx) = mpsc::channel(5);
@@ -378,7 +391,7 @@ mod tests {
         assert_eq!(tracker.len(), 0);
 
         // Simulate receiving a request by the [`NodeTask`]
-        recv_and_insert(&mut rx, &mut tracker).await;
+        recv_and_insert(&mut rx, &mut tracker, task_id).await;
 
         // We now have a request in the tracker
         assert_eq!(tracker.len(), 1);
@@ -401,7 +414,7 @@ mod tests {
         });
 
         // Simulate receiving a request by the [`NodeTask`]
-        recv_and_insert(&mut rx, &mut tracker).await;
+        recv_and_insert(&mut rx, &mut tracker, task_id).await;
         assert_eq!(tracker.len(), 2);
 
         // We still haven't actually completed any operations yet
@@ -456,6 +469,11 @@ mod tests {
         let proxy = Proxy::new(tx.clone());
         let mut tracker = Tracker::new();
 
+        // In real code, the `tokio::task::ID` is the id of the
+        // `EstablishedConnectionTask`. However, we are simulating those
+        // connections here, so just use an ID of an arbitrary task.
+        let task_id = task::spawn(async {}).id();
+
         let requests_completed = Arc::new(AtomicUsize::new(0));
 
         // This is the first "user" task that will issue proxy operations
@@ -473,7 +491,7 @@ mod tests {
         assert_eq!(tracker.len(), 0);
 
         // Simulate receiving a request by the [`NodeTask`]
-        recv_and_insert(&mut rx, &mut tracker).await;
+        recv_and_insert(&mut rx, &mut tracker, task_id).await;
 
         // We now have a request in the tracker
         assert_eq!(tracker.len(), 1);
@@ -517,6 +535,11 @@ mod tests {
         };
         let rack_id = RackUuid::new_v4();
 
+        // In real code, the `tokio::task::ID` is the id of the
+        // `EstablishedConnectionTask`. However, we are simulating those
+        // connections here, so just use an ID of an arbitrary task.
+        let task_id = task::spawn(async {}).id();
+
         // Test channel where the sender is usually cloned from the [`crate::NodeTaskHandle`],
         // and the receiver is owned by the local [`crate::NodeTask`]
         let (tx, mut rx) = mpsc::channel(5);
@@ -540,7 +563,7 @@ mod tests {
         assert_eq!(tracker.len(), 0);
 
         // Simulate receiving a request by the [`NodeTask`]
-        recv_and_insert(&mut rx, &mut tracker).await;
+        recv_and_insert(&mut rx, &mut tracker, task_id).await;
 
         // We now have a request in the tracker
         assert_eq!(tracker.len(), 1);
@@ -560,14 +583,15 @@ mod tests {
         });
 
         // Simulate receiving a request by the [`NodeTask`]
-        recv_and_insert(&mut rx, &mut tracker).await;
+        recv_and_insert(&mut rx, &mut tracker, task_id).await;
         assert_eq!(tracker.len(), 2);
 
         // We still haven't actually completed any operations yet
         assert_eq!(requests_completed.load(Ordering::Relaxed), 0);
 
-        // Now simulate a disconnection to the proxy destination
-        tracker.on_disconnect(&destination);
+        // Now simulate a disconnection to the proxy destination for this
+        // specific connection task.
+        tracker.on_disconnect(task_id);
 
         // Now wait for both responses to be processed
         wait_for_condition(