[HA] swbus: replace stale peer connection on reconnect

crates/swbus-core/src/mux/conn.rs

@@ -53,8 +53,14 @@ impl SwbusConn {
         SwbusConnProxy::new(self.send_queue_tx.clone())
     }
 
-    pub async fn shutdown(&self) -> Result<()> {
+    /// Signal the worker to stop without requiring an async shutdown path.
+    /// Connection replacement uses this after synchronously removing the old routes.
+    pub(crate) fn cancel(&self) {
         self.shutdown_ct.cancel();
+    }
+
+    pub async fn shutdown(&self) -> Result<()> {
+        self.cancel();
         Ok(())
     }
 }

crates/swbus-core/src/mux/conn_store.rs

@@ -4,6 +4,7 @@ use crate::mux::SwbusConnMode;
 use crate::mux::SwbusMultiplexer;
 use dashmap::DashMap;
 use std::sync::Arc;
+use std::sync::Mutex;
 use swbus_config::PeerConfig;
 use tokio::time::Duration;
 use tokio_util::sync::CancellationToken;
@@ -18,13 +19,16 @@ enum ConnTracker {
 pub struct SwbusConnStore {
     mux: Arc<SwbusMultiplexer>,
     connections: DashMap<String, ConnTracker>,
+    // Keep peer replacement and registration atomic across concurrent connection handshakes.
+    connection_update_lock: Mutex<()>,
 }
 
 impl SwbusConnStore {
     pub fn new(mux: Arc<SwbusMultiplexer>) -> Self {
         SwbusConnStore {
             mux,
             connections: DashMap::new(),
+            connection_update_lock: Mutex::new(()),
         }
     }
 
@@ -83,11 +87,58 @@ impl SwbusConnStore {
     }
 
     pub fn conn_established(&self, conn: SwbusConn) {
+        // A single critical section prevents two reconnects from each preserving the other as stale.
+        let _guard = self.connection_update_lock.lock().unwrap();
+        self.replace_stale_server_connections(conn.info());
         self.mux.register(conn.info(), conn.new_proxy());
         self.connections
             .insert(conn.info().id().to_string(), ConnTracker::SwbusConn(conn));
     }
 
+    /// Remove older accepted connections when the same logical peer reconnects.
+    ///
+    /// Server connection IDs contain the peer's ephemeral source port, so a reboot can create
+    /// a new ID while the old half-open connection remains routable. Match those connections by
+    /// peer service path and connection type. Client-mode connections are kept because they are
+    /// independently owned outbound paths with stable endpoint-based IDs.
+    fn replace_stale_server_connections(&self, new_conn_info: &SwbusConnInfo) {
+        if new_conn_info.mode() != SwbusConnMode::Server || new_conn_info.remote_service_path().is_none() {
+            return;
+        }
+
+        // Collect IDs first so no DashMap entry guard is held while removing connections or routes.
+        let stale_connection_ids: Vec<String> = self
+            .connections
+            .iter()
+            .filter_map(|entry| match entry.value() {
+                ConnTracker::SwbusConn(existing_conn)
+                    if existing_conn.info().id() != new_conn_info.id()
+                        && existing_conn.info().mode() == new_conn_info.mode()
+                        && existing_conn.info().connection_type() == new_conn_info.connection_type()
+                        && existing_conn.info().remote_service_path() == new_conn_info.remote_service_path() =>
+                {
+                    Some(entry.key().clone())
+                }
+                _ => None,
+            })
+            .collect();
+
+        for stale_connection_id in stale_connection_ids {
+            let Some((_, ConnTracker::SwbusConn(stale_conn))) = self.connections.remove(&stale_connection_id) else {
+                continue;
+            };
+            info!(
+                old_conn_id = stale_conn.info().id(),
+                new_conn_id = new_conn_info.id(),
+                peer = new_conn_info.remote_service_path().as_ref().unwrap().to_longest_path(),
+                "Replacing stale peer connection"
+            );
+            // Remove stale routes immediately; cancellation then lets the old worker exit cleanly.
+            self.mux.unregister(stale_conn.info());
+            stale_conn.cancel();
+        }
+    }
+
     pub async fn shutdown(&self) {
         for entry in self.connections.iter() {
             match entry.value() {
@@ -112,9 +163,7 @@ impl SwbusConnStore {
 mod tests {
     use super::*;
     use swbus_config::RouteConfig;
-    use swbus_proto::swbus::ConnectionType;
-    use swbus_proto::swbus::RouteScope;
-    use swbus_proto::swbus::ServicePath;
+    use swbus_proto::swbus::*;
     use tokio::sync::mpsc;
     #[tokio::test]
     async fn test_add_peer() {
@@ -177,4 +226,57 @@ mod tests {
             .iter()
             .any(|entry| entry.key() == conn_info.id() && matches!(entry.value(), ConnTracker::SwbusConn(_))));
     }
+
+    #[tokio::test]
+    async fn test_conn_established_replaces_stale_server_connection() {
+        let local_sp = ServicePath::from_string("region-a.cluster-a.10.0.0.1-dpu0").unwrap();
+        let remote_sp = ServicePath::from_string("region-a.cluster-a.10.0.0.2-dpu0").unwrap();
+        let mux = Arc::new(SwbusMultiplexer::new(vec![RouteConfig {
+            key: local_sp.clone(),
+            scope: RouteScope::InCluster,
+        }]));
+        let conn_store = SwbusConnStore::new(mux.clone());
+
+        let client_info = Arc::new(
+            SwbusConnInfo::new_client(ConnectionType::InCluster, "127.0.0.1:9000".parse().unwrap())
+                .with_remote_service_path(remote_sp.clone()),
+        );
+        let (client_tx, mut client_rx) = mpsc::channel(1);
+        conn_store.conn_established(SwbusConn::new(&client_info, client_tx));
+
+        let old_info = Arc::new(SwbusConnInfo::new_server(
+            ConnectionType::InCluster,
+            "127.0.0.1:8080".parse().unwrap(),
+            remote_sp.clone(),
+        ));
+        let (old_tx, mut old_rx) = mpsc::channel(1);
+        conn_store.conn_established(SwbusConn::new(&old_info, old_tx));
+
+        let new_info = Arc::new(SwbusConnInfo::new_server(
+            ConnectionType::InCluster,
+            "127.0.0.1:8081".parse().unwrap(),
+            remote_sp.clone(),
+        ));
+        let (new_tx, mut new_rx) = mpsc::channel(1);
+        conn_store.conn_established(SwbusConn::new(&new_info, new_tx));
+
+        assert!(!conn_store.connections.contains_key(old_info.id()));
+        assert!(conn_store.connections.contains_key(new_info.id()));
+        assert!(conn_store.connections.contains_key(client_info.id()));
+
+        let message = SwbusMessage {
+            header: Some(SwbusMessageHeader::new(local_sp, remote_sp, 1)),
+            body: Some(swbus_message::Body::PingRequest(PingRequest::new())),
+        };
+        mux.route_message(message.clone()).await.unwrap();
+
+        let mut expected = message;
+        expected.header.as_mut().unwrap().ttl -= 1;
+        assert_eq!(new_rx.recv().await.unwrap().unwrap(), expected);
+        assert!(matches!(
+            old_rx.try_recv(),
+            Err(mpsc::error::TryRecvError::Disconnected)
+        ));
+        assert!(matches!(client_rx.try_recv(), Err(mpsc::error::TryRecvError::Empty)));
+    }
 }

crates/swbus-core/src/mux/multiplexer.rs

@@ -252,17 +252,16 @@ impl SwbusMultiplexer {
         }
     }
 
-    // Update route for the give service path with the new nexthop. It returns a tuple of (nh_added, route_added).
-    // if inserted is true, the nexthop is newly inserted or updated (hop count changed).
-    // if route_added is true, the route entry is newly created.
+    // Add a next hop or refresh the proxy stored for an equal next-hop identity.
+    // Returns whether a new next hop and a new route entry were added, respectively.
     #[instrument(name = "update_route", level = "info", skip(self, nexthop), fields(nh_type=?nexthop.nh_type(), hop_count=nexthop.hop_count(), conn_info=nexthop.conn_info().as_ref().map(|x| x.id()).unwrap_or(&"None".to_string())))]
     fn update_route(&self, route_key: String, nexthop: SwbusNextHop) -> (bool, bool) {
         // If route entry doesn't exist, we insert the next hop as a new one.
         let mut route_added = false;
         if self.routes.get(&route_key).is_none() {
             route_added = true;
         }
-        let nh_added = self.routes.entry(route_key).or_default().insert(nexthop.clone());
+        let nh_added = self.routes.entry(route_key).or_default().replace(nexthop).is_none();
         info!("Update route entry: nh_added={nh_added}, route_added={route_added}");
         (nh_added, route_added)
     }
@@ -1448,6 +1447,40 @@ mod tests {
         assert!(matches!(error, SwbusError::InternalError { code: _, detail: _ }));
     }
 
+    #[tokio::test]
+    async fn test_update_route_replaces_existing_nexthop_proxy() {
+        let local_sp = ServicePath::from_string("region-a.cluster-a.10.0.0.2-dpu0").unwrap();
+        let remote_sp = ServicePath::from_string("region-a.cluster-a.10.0.0.1-dpu0").unwrap();
+        let mux = Arc::new(SwbusMultiplexer::new(vec![RouteConfig {
+            key: local_sp.clone(),
+            scope: RouteScope::InCluster,
+        }]));
+        let conn_info = Arc::new(SwbusConnInfo::new_server(
+            ConnectionType::InCluster,
+            "127.0.0.1:8080".parse().unwrap(),
+            remote_sp.clone(),
+        ));
+
+        let (old_tx, mut old_rx) = mpsc::channel(1);
+        let old_conn = SwbusConn::new(&conn_info, old_tx);
+        mux.register(old_conn.info(), old_conn.new_proxy());
+
+        let (new_tx, mut new_rx) = mpsc::channel(1);
+        let new_conn = SwbusConn::new(&conn_info, new_tx);
+        mux.register(new_conn.info(), new_conn.new_proxy());
+
+        let message = SwbusMessage {
+            header: Some(SwbusMessageHeader::new(local_sp, remote_sp, 1)),
+            body: Some(swbus_message::Body::PingRequest(PingRequest::new())),
+        };
+        mux.route_message(message.clone()).await.unwrap();
+
+        let mut expected = message;
+        expected.header.as_mut().unwrap().ttl -= 1;
+        assert_eq!(new_rx.recv().await.unwrap().unwrap(), expected);
+        assert!(matches!(old_rx.try_recv(), Err(mpsc::error::TryRecvError::Empty)));
+    }
+
     #[test]
     fn test_remove_existing_nexthop() {
         let route_config = RouteConfig {