53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
|
|
import json
|
|
import asyncio
|
|
import websockets
|
|
import threading
|
|
import queue
|
|
|
|
from utils.logger import get_logger
|
|
logger = get_logger(__name__)
|
|
|
|
# ========================= WebSocket 服务线程 =========================
|
|
class WebSocketSender(threading.Thread):
|
|
def __init__(self, send_queue: queue.Queue, stop_event: threading.Event, ws_host: str, ws_port: int):
|
|
super().__init__(daemon=True)
|
|
self.send_queue = send_queue
|
|
self.stop_event = stop_event
|
|
self.ws_clients = set()
|
|
self.ws_host = ws_host
|
|
self.ws_port = ws_port
|
|
|
|
async def _ws_handler(self, websocket):
|
|
self.ws_clients.add(websocket)
|
|
try:
|
|
async for _ in websocket:
|
|
pass
|
|
finally:
|
|
self.ws_clients.discard(websocket)
|
|
|
|
async def _broadcaster(self):
|
|
while not self.stop_event.is_set():
|
|
try:
|
|
msg = await asyncio.to_thread(self.send_queue.get, timeout=0.5)
|
|
except queue.Empty:
|
|
continue
|
|
data = json.dumps(msg)
|
|
dead = []
|
|
for ws in list(self.ws_clients):
|
|
try:
|
|
await ws.send(data)
|
|
except:
|
|
dead.append(ws)
|
|
for ws in dead:
|
|
self.ws_clients.discard(ws)
|
|
self.send_queue.task_done()
|
|
|
|
async def _run_async(self):
|
|
async with websockets.serve(self._ws_handler, self.ws_host, self.ws_port):
|
|
logger.info(f"[INFO] WebSocket server started at ws://{self.ws_host}:{self.ws_port}")
|
|
await self._broadcaster()
|
|
|
|
def run(self):
|
|
asyncio.run(self._run_async())
|