49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Small async pool for per-request ASR engines."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from typing import Callable, Generic, Optional, TypeVar
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
@dataclass
|
|
class _PoolState(Generic[T]):
|
|
queue: asyncio.Queue[T]
|
|
|
|
|
|
class LocalEnginePool(Generic[T]):
|
|
"""Fixed-size lazy engine pool backed by ``asyncio.Queue``."""
|
|
|
|
def __init__(self, size: int, factory: Callable[[], T]):
|
|
self._size = max(1, size)
|
|
self._factory = factory
|
|
self._state: Optional[_PoolState[T]] = None
|
|
self._init_lock = threading.Lock()
|
|
|
|
def _ensure_state(self) -> _PoolState[T]:
|
|
if self._state is not None:
|
|
return self._state
|
|
|
|
with self._init_lock:
|
|
if self._state is None:
|
|
self._state = _PoolState(queue=asyncio.Queue(maxsize=self._size))
|
|
for _ in range(self._size):
|
|
self._state.queue.put_nowait(self._factory())
|
|
return self._state
|
|
|
|
def warmup(self) -> None:
|
|
self._ensure_state()
|
|
|
|
async def acquire(self) -> T:
|
|
state = self._ensure_state()
|
|
return await state.queue.get()
|
|
|
|
async def release(self, engine: T) -> None:
|
|
state = self._ensure_state()
|
|
await state.queue.put(engine)
|