Kris
9a8a27bf36
本次提交新增了完整的多渠道消息网关系统，包括： 1. 支持飞书、钉钉、Web、Hook 四种渠道的适配器与配置 2. 领域模型层：消息、会话、绑定、出箱等核心实体 3. 应用服务层：管道、中间件、DTO 与业务逻辑 4. 基础设施层：持久化、过滤器、队列等端口实现 5. 接口层：REST API、SSE、WebSocket 通信端点 6. 前端页面与路由配置，添加渠道管理菜单 7. 新增相关依赖包与 docker-compose 部署配置
150 lines
4.6 KiB
Python
150 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from dataclasses import dataclass
|
|
|
|
from yuxi.channel.domain.port.queue_port import QueuePort
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class WorkerPoolConfig:
|
|
num_workers: int = 4
|
|
max_concurrent: int = 20
|
|
poll_timeout_ms: int = 5000
|
|
|
|
|
|
class WorkerPool:
|
|
def __init__(
|
|
self,
|
|
queue_port: QueuePort,
|
|
dispatch_fn,
|
|
*,
|
|
config: WorkerPoolConfig,
|
|
session_router: bool = True,
|
|
) -> None:
|
|
self._queue = queue_port
|
|
self._dispatch = dispatch_fn
|
|
self._config = config
|
|
self._session_router = session_router
|
|
self._running = False
|
|
self._tasks: list[asyncio.Task] = []
|
|
self._session_queues: list[asyncio.Queue] = []
|
|
self._rr_index = 0
|
|
|
|
async def start(self) -> None:
|
|
if hasattr(self._queue, "ensure_group"):
|
|
await self._queue.ensure_group()
|
|
self._running = True
|
|
|
|
for i in range(self._config.num_workers):
|
|
q: asyncio.Queue = asyncio.Queue(maxsize=self._config.max_concurrent)
|
|
self._session_queues.append(q)
|
|
task = asyncio.create_task(self._worker_loop(i, q))
|
|
self._tasks.append(task)
|
|
|
|
if self._session_router:
|
|
task = asyncio.create_task(self._dispatch_loop())
|
|
self._tasks.append(task)
|
|
|
|
logger.info(
|
|
"worker pool started: %d workers, session_router=%s",
|
|
self._config.num_workers,
|
|
self._session_router,
|
|
)
|
|
|
|
async def stop(self) -> None:
|
|
self._running = False
|
|
for task in self._tasks:
|
|
task.cancel()
|
|
await asyncio.gather(*self._tasks, return_exceptions=True)
|
|
self._tasks.clear()
|
|
logger.info("worker pool stopped")
|
|
|
|
@property
|
|
def is_running(self) -> bool:
|
|
return self._running
|
|
|
|
async def _dispatch_loop(self) -> None:
|
|
consumer_name = "dispatch-router"
|
|
|
|
while self._running:
|
|
try:
|
|
messages = await self._queue.dequeue(
|
|
count=self._config.max_concurrent,
|
|
block=self._config.poll_timeout_ms,
|
|
consumer_name=consumer_name,
|
|
)
|
|
if not messages:
|
|
continue
|
|
|
|
for msg in messages:
|
|
await self._route(msg)
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as exc:
|
|
logger.error("dispatch loop error: %s", exc)
|
|
await asyncio.sleep(1)
|
|
|
|
async def _route(self, msg: dict) -> None:
|
|
session_id = msg.get("session_id", "")
|
|
if session_id:
|
|
idx = hash(session_id) % self._config.num_workers
|
|
else:
|
|
idx = self._rr_index % self._config.num_workers
|
|
self._rr_index += 1
|
|
await self._session_queues[idx].put(msg)
|
|
|
|
async def _worker_loop(self, worker_id: int, queue: asyncio.Queue) -> None:
|
|
if not self._session_router:
|
|
await self._legacy_worker_loop(worker_id)
|
|
return
|
|
|
|
while self._running:
|
|
try:
|
|
msg = await asyncio.wait_for(queue.get(), timeout=1.0)
|
|
except TimeoutError:
|
|
continue
|
|
except asyncio.CancelledError:
|
|
break
|
|
|
|
try:
|
|
await self._dispatch(msg)
|
|
except Exception as exc:
|
|
logger.error("dispatch error: %s", exc)
|
|
continue
|
|
|
|
stream_id = msg.get("_stream_id")
|
|
if stream_id:
|
|
await self._queue.ack(stream_id)
|
|
|
|
async def _legacy_worker_loop(self, worker_id: int) -> None:
|
|
consumer_name = f"worker-{worker_id}"
|
|
|
|
while self._running:
|
|
try:
|
|
messages = await self._queue.dequeue(
|
|
count=self._config.max_concurrent,
|
|
block=self._config.poll_timeout_ms,
|
|
consumer_name=consumer_name,
|
|
)
|
|
if not messages:
|
|
continue
|
|
|
|
for msg in messages:
|
|
try:
|
|
await self._dispatch(msg)
|
|
except Exception as exc:
|
|
logger.error("dispatch error: %s", exc)
|
|
continue
|
|
stream_id = msg.get("_stream_id")
|
|
if stream_id:
|
|
await self._queue.ack(stream_id)
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as exc:
|
|
logger.error("worker loop error: %s", exc)
|
|
await asyncio.sleep(1)
|