134 lines
5.3 KiB
Python
134 lines
5.3 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 logging
|
|
import os
|
|
import time
|
|
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")
|
|
_import_limiter: RunVar[CapacityLimiter] = RunVar("hockey_import_admission")
|
|
logger = logging.getLogger("hockey.import")
|
|
|
|
|
|
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
|
|
|
|
|
|
async def run_import_db(function: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
|
|
"""Admit one bulk import before it borrows a shared DB worker.
|
|
|
|
RLock waiters remain asynchronous instead of consuming every worker needed
|
|
by control requests. The limiter belongs to the current event loop and is
|
|
shared by services. With at least two DB workers, an import leaves capacity
|
|
for ordinary DB work. Network I/O must happen before this boundary.
|
|
"""
|
|
if _atomic_session.get() is not None:
|
|
raise RuntimeError("An atomic database transaction cannot cross worker threads")
|
|
try:
|
|
limiter = _import_limiter.get()
|
|
except LookupError:
|
|
limiter = CapacityLimiter(1)
|
|
_import_limiter.set(limiter)
|
|
queued_at = time.monotonic()
|
|
# Do not detach run_db: it joins an already running worker on cancellation,
|
|
# so this admission token remains held until its transaction has finished.
|
|
async with limiter:
|
|
started_at = time.monotonic()
|
|
wait_ms = (started_at - queued_at) * 1000
|
|
function_name = getattr(function, "__qualname__", type(function).__name__)
|
|
details = {"function": function_name, "queue_wait_ms": round(wait_ms, 1)}
|
|
if wait_ms >= 250:
|
|
logger.info("Импорт дождался своей очереди", extra={
|
|
"category": "import", "event": "import.queue_wait", "details": details,
|
|
})
|
|
outcome = "completed"
|
|
try:
|
|
return await run_db(function, *args, **kwargs)
|
|
except asyncio.CancelledError:
|
|
outcome = "cancelled_after_worker_finished"
|
|
raise
|
|
except Exception as error:
|
|
outcome = "failed"
|
|
logger.error("Ошибка выполнения импорта", extra={
|
|
"category": "import", "event": "import.failed",
|
|
"details": {**details, "error_type": type(error).__name__},
|
|
})
|
|
raise
|
|
finally:
|
|
elapsed_ms = (time.monotonic() - started_at) * 1000
|
|
if wait_ms >= 250 or elapsed_ms >= 1000:
|
|
logger.info("Выполнение импорта завершено", extra={
|
|
"category": "import", "event": "import.finished",
|
|
"details": {**details, "elapsed_ms": round(elapsed_ms, 1), "outcome": outcome},
|
|
})
|
|
|
|
|
|
def _db_endpoint(function: Callable[..., Any], runner: Callable[..., Any]) -> Callable[..., Any]:
|
|
signature = inspect.signature(function, eval_str=True)
|
|
|
|
@wraps(function)
|
|
async def endpoint(*args: Any, **kwargs: Any) -> Any:
|
|
return await runner(function, *args, **kwargs)
|
|
|
|
endpoint.__signature__ = signature
|
|
return endpoint
|
|
|
|
|
|
def db_endpoint(function: Callable[..., Any]) -> Callable[..., Any]:
|
|
"""Run a synchronous FastAPI endpoint in the bounded DB worker pool."""
|
|
return _db_endpoint(function, run_db)
|
|
|
|
|
|
def import_db_endpoint(function: Callable[..., Any]) -> Callable[..., Any]:
|
|
"""Admit a synchronous bulk-import endpoint before dispatching its worker."""
|
|
return _db_endpoint(function, run_import_db)
|