53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
"""Bound synchronous database work without blocking the ASGI event loop.
|
|
|
|
Each job opens and closes its own Session. A transaction must stay inside one
|
|
job; passing an existing atomic Session to another thread is rejected.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
import os
|
|
from functools import partial, wraps
|
|
from typing import Any, Callable
|
|
|
|
from anyio import CapacityLimiter, to_thread
|
|
from anyio.lowlevel import RunVar
|
|
|
|
from .database import _atomic_session
|
|
|
|
_limiter: RunVar[CapacityLimiter] = RunVar("hockey_database_workers")
|
|
|
|
|
|
def worker_limit() -> int:
|
|
try:
|
|
configured = int(os.getenv("HOCKEY_DB_WORKERS", "8"))
|
|
pool_size = int(os.getenv("HOCKEY_DB_POOL_SIZE", "10"))
|
|
except ValueError:
|
|
return 8
|
|
return max(1, min(32, configured, max(1, pool_size)))
|
|
|
|
|
|
async def run_db(function: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
|
|
if _atomic_session.get() is not None:
|
|
raise RuntimeError("An atomic database transaction cannot cross worker threads")
|
|
try:
|
|
limiter = _limiter.get()
|
|
except LookupError:
|
|
limiter = CapacityLimiter(worker_limit())
|
|
_limiter.set(limiter)
|
|
# Do not abandon a worker on cancellation: its transaction must finish or
|
|
# roll back before the caller can release the concurrency budget.
|
|
return await to_thread.run_sync(partial(function, *args, **kwargs), limiter=limiter)
|
|
|
|
|
|
def db_endpoint(function: Callable[..., Any]) -> Callable[..., Any]:
|
|
"""Run a synchronous FastAPI endpoint in the same bounded DB worker pool."""
|
|
signature = inspect.signature(function, eval_str=True)
|
|
|
|
@wraps(function)
|
|
async def endpoint(*args: Any, **kwargs: Any) -> Any:
|
|
return await run_db(function, *args, **kwargs)
|
|
|
|
endpoint.__signature__ = signature
|
|
return endpoint
|