ForcePilot/backend/package/yuxi/channel/worker/pool.py

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)