ForcePilot/backend/package/yuxi/channels/adapters/feishu/sequential.py

from __future__ import annotations

import asyncio
from collections import defaultdict
from collections.abc import Coroutine
from typing import Any

_SEQUENTIAL_TIMEOUT_S = 300


class FeishuSequentialQueue:
    def __init__(self, timeout_s: float = _SEQUENTIAL_TIMEOUT_S):
        self._queues: dict[str, asyncio.Queue] = {}
        self._locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
        self._timeout_s = timeout_s

    async def acquire(self, key: str) -> None:
        lock = self._locks[key]
        try:
            await asyncio.wait_for(lock.acquire(), timeout=self._timeout_s)
        except TimeoutError:
            raise RuntimeError(f"[Sequential] Timeout acquiring lock for key '{key}'")

    def release(self, key: str) -> None:
        lock = self._locks.get(key)
        if lock and lock.locked():
            lock.release()

    async def run_sequential(self, key: str, coro: Coroutine[Any, Any, Any]) -> None:
        try:
            await asyncio.wait_for(self.acquire(key), timeout=self._timeout_s)
            try:
                await coro
            finally:
                self.release(key)
        except TimeoutError:
            raise RuntimeError(f"[Sequential] Timeout waiting for key '{key}'")

    def remove(self, key: str) -> None:
        lock = self._locks.pop(key, None)
        if lock and lock.locked():
            lock.release()

    def clear(self) -> None:
        for lock in self._locks.values():
            if lock.locked():
                lock.release()
        self._locks.clear()