balancer: check and handle closed channels more aggressively

Author: dyc3

crates/ott-balancer/src/balancer.rs

@@ -528,18 +528,31 @@ pub async fn join_monolith(
     let handle = tokio::task::Builder::new()
         .name(format!("monolith {}", monolith_id).as_ref())
         .spawn(async move {
-            while let Some(msg) = client_inbound_rx.recv().await {
+            loop {
+                let Some(msg) = (tokio::select! {
+                    biased;
+                    _ = monolith_outbound_tx.closed() => {
+                        None
+                    }
+                    msg = client_inbound_rx.recv() => {
+                        msg
+                    }
+                }) else {
+                    info!("monolith disconnected, stopping client inbound handler");
+                    break;
+                };
+
                 if let Err(e) =
                     handle_client_inbound(ctx.clone(), msg, monolith_outbound_tx.clone()).await
                 {
                     error!("failed to handle client inbound: {:?}", e);
                     if monolith_outbound_tx.is_closed() {
-                        // the monolith has disconnected
                         info!("monolith disconnected, stopping client inbound handler");
                         break;
                     }
                 }
             }
+            client_inbound_rx.close();
         })?;
     Ok(handle)
 }

crates/ott-balancer/src/client.rs

@@ -65,6 +65,9 @@ impl ClientLink {
     /// Receive the next message from the Balancer that needs to be sent to this client.
     pub async fn outbound_recv(&mut self) -> Result<SocketMessage, RecvError> {
         let msg = tokio::select! {
+            _ = self.room_tx.closed() => {
+                return Err(RecvError::Closed);
+            }
             msg = self.unicast_rx.recv() => {
                 match msg {
                     Some(msg) => Ok(msg),
@@ -197,8 +200,8 @@ pub async fn client_entry<'r>(
 
     loop {
         tokio::select! {
-            Ok(msg) = client_link.outbound_recv() => {
-                if let SocketMessage::Message(msg) = msg {
+            msg = client_link.outbound_recv() => {
+                if let Ok(SocketMessage::Message(msg)) = msg {
                     debug!(event = "ws", node_id = %client_id, room = %room_name, direction = "tx");
                     if let Err(err) = stream.send(msg).await {
                         error!("Error sending ws message to client: {:?}", err);
@@ -239,15 +242,24 @@ pub async fn client_entry<'r>(
     }
 
     info!("ending client connection");
-    client_link
-        .room_tx
-        .send(Context::new(
-            client_id,
-            SocketMessage::Message(Message::Close(Some(CloseFrame {
-                code: CloseCode::Normal,
-                reason: "client connection ended".into(),
-            }))),
-        ))
+    if !client_link.room_tx.is_closed() {
+        client_link
+            .room_tx
+            .send(Context::new(
+                client_id,
+                SocketMessage::Message(Message::Close(Some(CloseFrame {
+                    code: CloseCode::Normal,
+                    reason: "client connection ended".into(),
+                }))),
+            ))
+            .await?;
+    }
+
+    stream
+        .close(Some(CloseFrame {
+            code: CloseCode::Normal,
+            reason: "client connection ended".into(),
+        }))
         .await?;
 
     Ok(())

crates/ott-balancer/src/connection.rs

@@ -263,6 +263,10 @@ async fn connect_and_maintain(
                 }
             }
         }
+        // drain the outbound queue
+        outbound_rx.close();
+        while let Some(_) = outbound_rx.recv().await {}
+
         if cancel.is_cancelled() {
             break;
         }