72 lines
2.5 KiB
Python
72 lines
2.5 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 asyncio
|
|
import inspect
|
|
import os
|
|
from functools import partial, wraps
|
|
from typing import Any, Callable
|
|
|
|
from anyio import CancelScope, 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)
|
|
# asyncio.Task.cancel() bypasses AnyIO's cancel-scope protection. Keep the
|
|
# worker's await alive so its transaction and limiter token finish together.
|
|
worker = asyncio.create_task(
|
|
to_thread.run_sync(partial(function, *args, **kwargs), limiter=limiter)
|
|
)
|
|
try:
|
|
return await asyncio.shield(worker)
|
|
except asyncio.CancelledError:
|
|
# AnyIO scopes repeat cancellation at every await. Shield that scope
|
|
# too, otherwise cleanup would spin while the SQL thread is running.
|
|
with CancelScope(shield=True):
|
|
while not worker.done():
|
|
try:
|
|
await asyncio.shield(worker)
|
|
except asyncio.CancelledError:
|
|
continue
|
|
except Exception:
|
|
break
|
|
if not worker.cancelled():
|
|
worker.exception() # Consume errors while preserving cancellation.
|
|
raise
|
|
|
|
|
|
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
|