app/services/media_queue.py - master - artyuhovs/AIFeedBackTrainingBot

app/services/media_queue.py
108 строк4 KB
Codex
Добавить прозрачный медиаконвейер в версии 0.8.1
19 июл 2026, 20:24
19 июл 2026, 20:249c5c8db
from __future__ import annotations

import asyncio
import logging
from collections.abc import AsyncIterator, Awaitable, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class StageQueuePosition:
    active: bool
    position: int | None = None


PositionCallback = Callable[[StageQueuePosition], Awaitable[None]]


@dataclass(eq=False)
class _QueueEntry:
    ready: asyncio.Event
    on_position: PositionCallback | None
    active: bool = False
    last_position: StageQueuePosition | None = None
    notification_lock: asyncio.Lock = field(default_factory=asyncio.Lock)


class MediaProcessingQueue:
    """A bounded process-local FIFO stage with visible waiting positions."""

    def __init__(self, concurrency: int = 1) -> None:
        if concurrency < 1:
            raise ValueError("concurrency must be at least 1")
        self._concurrency = concurrency
        self._lock = asyncio.Lock()
        self._entries: list[_QueueEntry] = []

    @asynccontextmanager
    async def acquire(self, on_position: PositionCallback | None = None) -> AsyncIterator[None]:
        entry = _QueueEntry(asyncio.Event(), on_position)
        await self._enqueue(entry)
        try:
            await self._publish_position(entry)
            await entry.ready.wait()
            await self._publish_position(entry)
            yield
        finally:
            await self._release(entry)

    async def _enqueue(self, entry: _QueueEntry) -> None:
        async with self._lock:
            self._entries.append(entry)
            self._promote_locked()

    async def _release(self, entry: _QueueEntry) -> None:
        async with self._lock:
            try:
                index = self._entries.index(entry)
            except ValueError:
                return
            was_active = entry.active
            self._entries.pop(index)
            if was_active:
                self._promote_locked()
            waiting_entries = tuple(candidate for candidate in self._entries if not candidate.active)

        if waiting_entries:
            await asyncio.gather(*(self._publish_position(waiter) for waiter in waiting_entries))

    def _promote_locked(self) -> None:
        available = self._concurrency - sum(entry.active for entry in self._entries)
        if available <= 0:
            return
        for entry in self._entries:
            if entry.active:
                continue
            entry.active = True
            entry.ready.set()
            available -= 1
            if available == 0:
                return

    async def _publish_position(self, entry: _QueueEntry) -> StageQueuePosition | None:
        if entry.on_position is None:
            return None
        async with entry.notification_lock:
            async with self._lock:
                if entry not in self._entries:
                    return None
                waiting_entries = [candidate for candidate in self._entries if not candidate.active]
                waiting_position = None
                if not entry.active:
                    waiting_position = waiting_entries.index(entry) + 1
                position = StageQueuePosition(active=entry.active, position=waiting_position)
                if position == entry.last_position:
                    return position
                entry.last_position = position
            try:
                await entry.on_position(position)
            except Exception:
                logger.warning(
                    "media_queue.position_callback_failed",
                    extra={"operation": "media_queue.position_callback_failed"},
                    exc_info=True,
                )
            return position