refactor close and send methods of WebSocketStreamingServer to separate functionality

Author: janaab11

src/diart/websockets.py

@@ -94,16 +94,15 @@ def _create_client_state(self, client_id: Text) -> ClientState:
         pipeline = self.pipeline_class(self.pipeline_config)
 
         audio_source = src.WebSocketAudioSource(
-            uri=f"{self.uri}:{client_id}",
-            sample_rate=self.pipeline_config.sample_rate,
+            uri=f"{self.uri}:{client_id}", sample_rate=self.pipeline_config.sample_rate,
         )
 
         inference = StreamingInference(
             pipeline=pipeline,
             source=audio_source,
             # The following variables are fixed for a client
             batch_size=1,
-            do_profile=False, # for minimal latency
+            do_profile=False,  # for minimal latency
             do_plot=False,
             show_progress=False,
             progress_bar=None,
@@ -143,6 +142,11 @@ def _on_connect(self, client: Dict[Text, Any], server: WebsocketServer) -> None:
             self.send(client_id, "READY")
         except Exception as e:
             logger.error(f"Failed to initialize client {client_id}: {e}")
+
+            # Send close notification to client
+            self.send(client_id, "CLOSE")
+
+            # Close audio source and remove client
             self.close(client_id)
 
     def _on_disconnect(self, client: Dict[Text, Any], server: WebsocketServer) -> None:
@@ -157,6 +161,11 @@ def _on_disconnect(self, client: Dict[Text, Any], server: WebsocketServer) -> No
         """
         client_id = client["id"]
         logger.info(f"Client disconnected: {client_id}")
+
+        # Send close notification to client
+        self.send(client_id, "CLOSE")
+
+        # Close audio source and remove client
         self.close(client_id)
 
     def _on_message_received(
@@ -182,7 +191,13 @@ def _on_message_received(
             self._clients[client_id].audio_source.process_message(message)
         except (socket.error, ConnectionError) as e:
             logger.warning(f"Client {client_id} disconnected: {e}")
+
+            # Send close notification to client
+            self.send(client_id, "CLOSE")
+
+            # Close audio source and remove client
             self.close(client_id)
+
         except Exception as e:
             logger.error(f"Error processing message from client {client_id}: {e}")
             # Don't close the connection for non-connection related errors
@@ -202,45 +217,14 @@ def send(self, client_id: Text, message: AnyStr) -> None:
             return
 
         client = next((c for c in self.server.clients if c["id"] == client_id), None)
-
         if client is None:
             return
 
         try:
             self.server.send_message(client, message)
-        except (socket.error, ConnectionError) as e:
-            logger.warning(
-                f"Client {client_id} disconnected while sending message: {e}"
-            )
-            self.close(client_id)
         except Exception as e:
             logger.error(f"Failed to send message to client {client_id}: {e}")
 
-    def run(self) -> None:
-        """Start the WebSocket server."""
-        logger.info(f"Starting WebSocket server on {self.uri}")
-        max_retries = 3
-        retry_count = 0
-
-        while retry_count < max_retries:
-            try:
-                self.server.run_forever()
-                break  # If server exits normally, break the retry loop
-            except (socket.error, ConnectionError) as e:
-                logger.warning(f"WebSocket connection error: {e}")
-                retry_count += 1
-                if retry_count < max_retries:
-                    logger.info(
-                        f"Attempting to restart server (attempt {retry_count + 1}/{max_retries})"
-                    )
-                else:
-                    logger.error("Max retry attempts reached. Server shutting down.")
-            except Exception as e:
-                logger.error(f"Fatal server error: {e}")
-                break
-            finally:
-                self.close_all()
-
     def close(self, client_id: Text) -> None:
         """Close a specific client's connection and cleanup resources.
 
@@ -261,22 +245,6 @@ def close(self, client_id: Text) -> None:
             client_state.audio_source.close()
             del self._clients[client_id]
 
-            # Try to send a close frame to the client
-            client = next((c for c in self.server.clients if c["id"] == client_id), None)
-
-            if client is None:
-                return
-
-            try:
-                self.server.send_message(client, "CLOSE")
-            except (socket.error, ConnectionError) as e:
-                logger.warning(
-                    f"Client {client_id} disconnected while sending message: {e}"
-                )
-                self.close(client_id)
-            except Exception as e:
-                logger.error(f"Failed to send message to client {client_id}: {e}")
-
             logger.info(
                 f"Closed connection and cleaned up state for client: {client_id}"
             )
@@ -290,9 +258,40 @@ def close_all(self) -> None:
         logger.info("Shutting down server...")
         try:
             for client_id in self._clients.keys():
+                # Close audio source and remove client
                 self.close(client_id)
+
+                # Send close notification to client
+                self.send(client_id, "CLOSE")
+
             if self.server is not None:
                 self.server.shutdown_gracefully()
+
             logger.info("Server shutdown complete")
         except Exception as e:
             logger.error(f"Error during shutdown: {e}")
+
+    def run(self) -> None:
+        """Start the WebSocket server."""
+        logger.info(f"Starting WebSocket server on {self.uri}")
+        max_retries = 3
+        retry_count = 0
+
+        while retry_count < max_retries:
+            try:
+                self.server.run_forever()
+                break  # If server exits normally, break the retry loop
+            except (socket.error, ConnectionError) as e:
+                logger.warning(f"WebSocket connection error: {e}")
+                retry_count += 1
+                if retry_count < max_retries:
+                    logger.info(
+                        f"Attempting to restart server (attempt {retry_count + 1}/{max_retries})"
+                    )
+                else:
+                    logger.error("Max retry attempts reached. Server shutting down.")
+            except Exception as e:
+                logger.error(f"Fatal server error: {e}")
+                break
+            finally:
+                self.close_all()