mirror of
https://github.com/yogeshojha/rengine.git
synced 2026-10-20 15:21:30 +02:00
Yogesh Ojha
ad47eb4fb9
Stage tasks hold a slot for their whole tool run, so step hand-overs and finalize queued behind them and background jobs took scan slots. - worker-scans consumes scans (run_scan_stage only), scalable with --scale - worker-control consumes scan_control: launch, hand-over, resume, finalize - worker-default consumes default and critical - reaper reads both scan queues; readiness requires a consumer on every queue - connection budget counts every worker; max_connections 250
84 lines
2.4 KiB
Python
84 lines
2.4 KiB
Python
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
from app.config import settings
|
|
from shared.logging import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
_server_settings = {"application_name": f"{settings.APP_NAME}-api"}
|
|
|
|
if settings.DB_IDLE_TX_TIMEOUT > 0:
|
|
_server_settings["idle_in_transaction_session_timeout"] = str(
|
|
settings.DB_IDLE_TX_TIMEOUT * 1000
|
|
)
|
|
|
|
engine = create_async_engine(
|
|
settings.database_url,
|
|
echo=settings.SQL_ECHO,
|
|
future=True,
|
|
pool_size=settings.DB_POOL_SIZE,
|
|
max_overflow=settings.DB_MAX_OVERFLOW,
|
|
pool_timeout=settings.DB_POOL_TIMEOUT,
|
|
pool_recycle=settings.DB_POOL_RECYCLE,
|
|
pool_pre_ping=True,
|
|
connect_args={"server_settings": _server_settings},
|
|
)
|
|
|
|
async_db_session = async_sessionmaker(
|
|
engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
)
|
|
|
|
|
|
async def get_session() -> AsyncSession:
|
|
async with async_db_session() as session:
|
|
yield session
|
|
|
|
|
|
def pool_stats() -> dict[str, int]:
|
|
pool = engine.pool
|
|
return {
|
|
"size": pool.size(),
|
|
"checked_out": pool.checkedout(),
|
|
"checked_in": pool.checkedin(),
|
|
"overflow": pool.overflow(),
|
|
"capacity": settings.DB_POOL_SIZE + settings.DB_MAX_OVERFLOW,
|
|
}
|
|
|
|
|
|
def pool_demand() -> int:
|
|
"""Connections every pool can ask for at once."""
|
|
per_child = settings.WORKER_DB_POOL_SIZE + settings.WORKER_DB_MAX_OVERFLOW
|
|
return (
|
|
settings.DB_POOL_SIZE
|
|
+ settings.DB_MAX_OVERFLOW
|
|
+ settings.worker_children * per_child
|
|
+ per_child
|
|
)
|
|
|
|
|
|
async def check_capacity() -> None:
|
|
"""Warn when pool demand exceeds max_connections."""
|
|
demand = pool_demand()
|
|
try:
|
|
async with engine.connect() as conn:
|
|
allowed = int(await conn.scalar(text("SHOW max_connections")))
|
|
reserved = int(
|
|
await conn.scalar(text("SHOW superuser_reserved_connections"))
|
|
)
|
|
except Exception:
|
|
logger.debug("could not read the connection ceiling", exc_info=True)
|
|
return
|
|
usable = allowed - reserved
|
|
if demand > usable:
|
|
logger.warning(
|
|
"database pools can outgrow the server",
|
|
demand=demand,
|
|
usable=usable,
|
|
hint="raise POSTGRES_MAX_CONNECTIONS or lower the pool sizes",
|
|
)
|
|
else:
|
|
logger.info("database pool headroom", demand=demand, usable=usable)
|