Merge pull request #504 from suofang/fix_not_clear_threads

fix:解决客户端websocket连接断开后异步消费线程不销毁的BUG

main/xiaozhi-server/core/connection.py

@@ -150,31 +150,29 @@ class ConnectionHandler:
             # 异步初始化
             self.executor.submit(self._initialize_components)
             # tts 消化线程
-            tts_priority = threading.Thread(target=self._tts_priority_thread, daemon=True)
-            tts_priority.start()
+            self.tts_priority_thread = threading.Thread(target=self._tts_priority_thread, daemon=True)
+            self.tts_priority_thread.start()
 
             # 音频播放 消化线程
-            audio_play_priority = threading.Thread(target=self._audio_play_priority_thread, daemon=True)
-            audio_play_priority.start()
+            self.audio_play_priority_thread = threading.Thread(target=self._audio_play_priority_thread, daemon=True)
+            self.audio_play_priority_thread.start()
 
             try:
                 async for message in self.websocket:
                     await self._route_message(message)
             except websockets.exceptions.ConnectionClosed:
                 self.logger.bind(tag=TAG).info("客户端断开连接")
-                await self.close()
 
         except AuthenticationError as e:
             self.logger.bind(tag=TAG).error(f"Authentication failed: {str(e)}")
-            await ws.close()
             return
         except Exception as e:
             stack_trace = traceback.format_exc()
             self.logger.bind(tag=TAG).error(f"Connection error: {str(e)}-{stack_trace}")
-            await ws.close()
             return
         finally:
             await self.memory.save_memory(self.dialogue.dialogue)
+            await self.close(ws)
 
     async def _route_message(self, message):
         """消息路由"""
@@ -501,7 +499,12 @@ class ConnectionHandler:
         while not self.stop_event.is_set():
             text = None
             try:
-                future = self.tts_queue.get()
+                try:
+                    future = self.tts_queue.get(timeout=1)
+                except queue.Empty:
+                    if self.stop_event.is_set():
+                        break
+                    continue
                 if future is None:
                     continue
                 text = None
@@ -542,7 +545,12 @@ class ConnectionHandler:
         while not self.stop_event.is_set():
             text = None
             try:
-                opus_datas, text, text_index = self.audio_play_queue.get()
+                try:
+                    opus_datas, text, text_index = self.audio_play_queue.get(timeout=1)
+                except queue.Empty:
+                    if self.stop_event.is_set():
+                        break
+                    continue
                 future = asyncio.run_coroutine_threadsafe(sendAudioMessage(self, opus_datas, text, text_index),
                                                           self.loop)
                 future.result()
@@ -572,16 +580,41 @@ class ConnectionHandler:
             self.tts_first_text_index = text_index
         self.tts_last_text_index = text_index
 
-    async def close(self):
+    async def close(self, ws=None):
         """资源清理方法"""
 
-        # 清理其他资源
-        self.stop_event.set()
-        self.executor.shutdown(wait=False)
-        if self.websocket:
+        # 触发停止事件并清理资源
+        if self.stop_event:
+            self.stop_event.set()
+        
+        # 立即关闭线程池
+        if self.executor:
+            self.executor.shutdown(wait=False, cancel_futures=True)
+            self.executor = None
+        
+        # 清空任务队列
+        self._clear_queues()
+        
+        if ws:
+            await ws.close()
+        elif self.websocket:
             await self.websocket.close()
         self.logger.bind(tag=TAG).info("连接资源已释放")
 
+    def _clear_queues(self):
+        # 清空所有任务队列
+        for q in [self.tts_queue, self.audio_play_queue]:
+            if not q:
+                continue
+            while not q.empty():
+                try:
+                    q.get_nowait()
+                except queue.Empty:
+                    continue
+            q.queue.clear()
+            # 添加毒丸信号到队列，确保线程退出
+            # q.queue.put(None)
+
     def reset_vad_states(self):
         self.client_audio_buffer = bytes()
         self.client_have_voice = False