如何正确销毁线程中运行的异步WebSocket流处理类实例

要实现该WebSocket类的优雅销毁，核心逻辑是主动关闭WebSocket连接打断消息监听的阻塞状态，同时取消心跳任务、回收异步资源和工作线程，具体修改如下：

核心修改逻辑 #

• 存储异步事件循环、心跳任务、主循环任务的引用，方便后续销毁时操作
• 主动调用ws.close()关闭连接，即可直接终止async for msg in self.ws的阻塞等待，无需额外修改消息循环逻辑
• 提供同步的close()方法，支持外部线程安全触发销毁流程
• 处理任务取消、连接关闭的异常，保证资源正常回收

修改后完整代码 #

import json
import aiohttp
import asyncio
import gzip
from threading import Thread
from typing import Optional

class WebSocket:
    KEEPALIVE_INTERVAL_S = 10

    def __init__(self, url, on_connect, on_msg):
        self.url = url
        self.on_connect = on_connect
        self.on_msg = on_msg

        self.streams = {}
        self._running = True
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._ping_task: Optional[asyncio.Task] = None
        self._main_task: Optional[asyncio.Task] = None
        self._ws: Optional[aiohttp.ClientWebSocketResponse] = None
        self.worker_thread = Thread(name='WebSocket', target=self.thread_func, daemon=True)
        self.worker_thread.start()

    def thread_func(self):
        asyncio.run(self.aio_run())

    async def aio_run(self):
        self._loop = asyncio.get_running_loop()
        try:
            async with aiohttp.ClientSession() as session:
                self._ws = await session.ws_connect(self.url)
                await self.on_connect(self)

                async def ping():
                    while self._running:
                        try:
                            print('KEEPALIVE')
                            await self._ws.ping()
                            await asyncio.sleep(WebSocket.KEEPALIVE_INTERVAL_S)
                        except asyncio.CancelledError:
                            break
                        except Exception as e:
                            print(f"心跳异常: {e}")
                            break

                async def main_loop():
                    try:
                        async for msg in self._ws:
                            def extract_data(msg):
                                if msg.type == aiohttp.WSMsgType.BINARY:
                                    as_bytes = gzip.decompress(msg.data)
                                    as_string = as_bytes.decode('utf8')
                                    as_json = json.loads(as_string)
                                    return as_json
                                elif msg.type == aiohttp.WSMsgType.TEXT:
                                    return json.loads(msg.data)
                                elif msg.type == aiohttp.WSMsgType.ERROR:
                                    print('⛔️ aiohttp.WSMsgType.ERROR')
                                return msg.data

                            data = extract_data(msg)
                            self.on_msg(data)
                    except Exception as e:
                        if not self._running:
                            # 主动关闭触发的异常属于正常流程，无需报错
                            return
                        print(f"消息循环异常: {e}")

                self._ping_task = asyncio.create_task(ping())
                self._main_task = asyncio.create_task(main_loop())

                await asyncio.gather(self._ping_task, self._main_task)
        finally:
            self._running = False

    async def send_json(self, J):
        if self._ws and not self._ws.closed:
            await self._ws.send_json(J)

    def close(self):
        """外部调用的同步销毁方法，线程安全"""
        if not self._running:
            return
        self._running = False

        # 线程安全的触发异步销毁逻辑
        def _async_stop():
            async def __stop():
                # 关闭WebSocket连接，直接终止消息循环的阻塞等待
                if self._ws and not self._ws.closed:
                    await self._ws.close()
                # 取消两个任务
                if self._ping_task and not self._ping_task.done():
                    self._ping_task.cancel()
                    try:
                        await self._ping_task
                    except asyncio.CancelledError:
                        pass
                if self._main_task and not self._main_task.done():
                    self._main_task.cancel()
                    try:
                        await self._main_task
                    except asyncio.CancelledError:
                        pass
            asyncio.create_task(__stop())

        if self._loop and self._loop.is_running():
            self._loop.call_soon_threadsafe(_async_stop)
        # 等待工作线程退出
        self.worker_thread.join(timeout=5)

调用方式 #

需要销毁实例时，直接调用实例的close()方法即可：

ws = WebSocket("wss://xxx", on_connect_callback, on_msg_callback)
# 使用完成后销毁
ws.close()

原理说明 #

消息循环的async for msg in self._ws阻塞状态不需要手动打断，只要调用self._ws.close()主动关闭WebSocket连接，aiohttp内部会终止迭代器，消息循环任务会自然退出。心跳任务通过CancelledError和_running标志双重控制退出，所有异步资源回收完成后工作线程会自动结束。

内容的提问来源于stack exchange，提问作者P i