Deriver Jitter (#765)
* fix(deriver): Remove connection retry logic and add jitter to polling interval * chore(docs): Update changelog and document new configurations * chore: increment version numbers
This commit is contained in:
Vineeth Voruganti
GitHub
parent
bb6dad9157
commit
9f26fdd2ea
|
|
@ -44,18 +44,15 @@ DB_CONNECTION_URI=postgresql+psycopg://postgres:postgres@localhost:5432/postgres
|
|||
# DB_POOL_CLASS=default
|
||||
# DB_POOL_SIZE=10
|
||||
# DB_MAX_OVERFLOW=20
|
||||
# DB_POOL_TIMEOUT=5 # seconds; must stay under DB_CONNECTION_RETRY_MAX_DELAY_SECONDS for a pooled (non-null) DB_POOL_CLASS
|
||||
# DB_POOL_TIMEOUT=5 # seconds a pooled checkout waits for a free connection (QueuePool only)
|
||||
# DB_POOL_RECYCLE=300
|
||||
# DB_POOL_PRE_PING=true
|
||||
# DB_POOL_USE_LIFO=true
|
||||
# DB_SQL_DEBUG=false
|
||||
# DB_TRACING=false
|
||||
# Bounded exponential-backoff retry around connection checkout (guards against
|
||||
# transient transaction-pooler saturation). Applied lazily on first DB use.
|
||||
# DB_CONNECTION_RETRY_ENABLED=true
|
||||
# DB_CONNECTION_RETRY_MAX_DELAY_SECONDS=10.0
|
||||
# DB_CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS=0.1
|
||||
# DB_CONNECTION_RETRY_BACKOFF_MAX_SECONDS=2.0
|
||||
# Per-connection establish timeout (seconds) so a single connection attempt
|
||||
# fails fast instead of hanging when the server/pooler is unreachable.
|
||||
# DB_CONNECT_TIMEOUT_SECONDS=2
|
||||
|
||||
# =============================================================================
|
||||
# Authentication Settings
|
||||
|
|
@ -115,6 +112,11 @@ LLM_OPENAI_API_KEY=your-api-key-here
|
|||
# DERIVER_POLLING_BACKOFF_ENABLED=true
|
||||
# DERIVER_POLLING_SLEEP_MAX_INTERVAL_SECONDS=30.0
|
||||
# DERIVER_POLLING_BACKOFF_MULTIPLIER=2.0
|
||||
# Jitter so instances that start together don't poll in lockstep. Startup: sleep
|
||||
# a random delay in [0, value] before the first poll (0.0 disables). Per-cycle:
|
||||
# multiply every poll sleep by a random factor in [1-ratio, 1+ratio] (0.0 disables).
|
||||
# DERIVER_POLLING_STARTUP_JITTER_SECONDS=30.0
|
||||
# DERIVER_POLLING_JITTER_RATIO=0.5
|
||||
# DERIVER_STALE_SESSION_TIMEOUT_MINUTES=5
|
||||
# DERIVER_QUEUE_ERROR_RETENTION_SECONDS=2592000 # 30 days
|
||||
# DERIVER_MODEL_CONFIG__TEMPERATURE=
|
||||
|
|
|
|||
14
CHANGELOG.md
14
CHANGELOG.md
|
|
@ -5,6 +5,20 @@ All notable changes to this project will be documented in this file.
|
|||
The format is based on [Keep a Changelog](http://keepachangelog.com/)
|
||||
and this project adheres to [Semantic Versioning](http://semver.org/).
|
||||
|
||||
## [3.0.9] - 2026-06-02
|
||||
|
||||
### Changed
|
||||
|
||||
- Connection acquisition is now a single attempt with no server-side retry, on a vanilla `AsyncSession`. A new `DB_CONNECT_TIMEOUT_SECONDS` (default 2s) bounds the attempt so a saturated or unreachable pooler fails fast instead of holding a client connection open to re-knock. A saturated DB now surfaces to the caller — the API returns an error and the deriver backs off and retries on a later poll — which lets the pooler drain rather than amplifying saturation.
|
||||
|
||||
### Added
|
||||
|
||||
- Deriver poll jitter so instances that start together don't poll in lockstep: `DERIVER_POLLING_STARTUP_JITTER_SECONDS` (random delay before the first poll, default 30s) and `DERIVER_POLLING_JITTER_RATIO` (±fraction applied to every poll sleep, default 0.5). Both disable at `0.0`; the underlying backoff schedule is unchanged.
|
||||
|
||||
### Removed
|
||||
|
||||
- Reverted the connection-checkout retry and `HonchoAsyncSession` custom session introduced in 3.0.8. Removed the `DB_CONNECTION_RETRY_ENABLED` / `DB_CONNECTION_RETRY_MAX_DELAY_SECONDS` / `DB_CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS` / `DB_CONNECTION_RETRY_BACKOFF_MAX_SECONDS` settings, the `db_connection_acquisitions{outcome=...}` Prometheus counter, and the `db.pool.acquire` Sentry span. Alerting built on `db_connection_acquisitions` should migrate to `db_pool_connections` / `db_queries_in_flight`.
|
||||
|
||||
## [3.0.8] - 2026-06-01
|
||||
|
||||
### Added
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@
|
|||
|
||||
---
|
||||
|
||||
|
||||
|
||||
[](https://pypi.org/project/honcho-ai/)
|
||||
[](https://npmjs.org/package/@honcho-ai/sdk)
|
||||
[](https://discord.gg/honcho)
|
||||
|
|
|
|||
|
|
@ -26,18 +26,14 @@ POOL_CLASS = "default"
|
|||
POOL_PRE_PING = true
|
||||
POOL_SIZE = 10
|
||||
MAX_OVERFLOW = 20
|
||||
POOL_TIMEOUT = 5 # seconds; must stay under CONNECTION_RETRY_MAX_DELAY_SECONDS
|
||||
POOL_TIMEOUT = 5 # seconds a pooled checkout waits for a free connection (QueuePool only)
|
||||
POOL_RECYCLE = 300 # seconds
|
||||
POOL_USE_LIFO = true
|
||||
SQL_DEBUG = false
|
||||
TRACING = false
|
||||
# Bounded retry around connection checkout (used by background/tracked_db and
|
||||
# the lazy request session). For a pooled (non-null) POOL_CLASS, POOL_TIMEOUT
|
||||
# must be < CONNECTION_RETRY_MAX_DELAY_SECONDS or startup validation fails.
|
||||
CONNECTION_RETRY_ENABLED = true
|
||||
CONNECTION_RETRY_MAX_DELAY_SECONDS = 10.0
|
||||
CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS = 0.1
|
||||
CONNECTION_RETRY_BACKOFF_MAX_SECONDS = 2.0
|
||||
# Per-connection establish timeout (seconds) so a single connection attempt
|
||||
# fails fast instead of hanging when the server/pooler is unreachable.
|
||||
CONNECT_TIMEOUT_SECONDS = 2
|
||||
|
||||
# Authentication settings
|
||||
[auth]
|
||||
|
|
@ -95,6 +91,12 @@ POLLING_SLEEP_INTERVAL_SECONDS = 1.0
|
|||
POLLING_BACKOFF_ENABLED = true
|
||||
POLLING_SLEEP_MAX_INTERVAL_SECONDS = 30.0
|
||||
POLLING_BACKOFF_MULTIPLIER = 2.0
|
||||
# Jitter so instances that start together don't poll in lockstep. Startup:
|
||||
# sleep a random delay in [0, POLLING_STARTUP_JITTER_SECONDS] before the first
|
||||
# poll (0.0 disables). Per-cycle: multiply every poll sleep by a random factor
|
||||
# in [1 - ratio, 1 + ratio] (0.5 -> [0.5x, 1.5x]; 0.0 disables).
|
||||
POLLING_STARTUP_JITTER_SECONDS = 30.0
|
||||
POLLING_JITTER_RATIO = 0.5
|
||||
STALE_SESSION_TIMEOUT_MINUTES = 5
|
||||
# QUEUE_ERROR_RETENTION_SECONDS = 2592000 # 30 days
|
||||
DEDUPLICATE = true
|
||||
|
|
|
|||
|
|
@ -30,7 +30,8 @@ This guide helps you match the right SDK version to your Honcho API version. New
|
|||
|
||||
| Honcho API Version | TypeScript SDK | Python SDK |
|
||||
|-------------------|---------------|------------|
|
||||
| v3.0.8 (Current) | v2.1.2 | v2.1.2 |
|
||||
| v3.0.9 (Current) | v2.1.2 | v2.1.2 |
|
||||
| v3.0.8 | v2.1.2 | v2.1.2 |
|
||||
| v3.0.7 | v2.1.2 | v2.1.2 |
|
||||
| v3.0.6 | v2.1.1 | v2.1.1 |
|
||||
| v3.0.5 | v2.1.0 | v2.1.0 |
|
||||
|
|
|
|||
|
|
@ -27,7 +27,21 @@ Welcome to the Honcho changelog! This section documents all notable changes to t
|
|||
### Honcho API and SDK Changelogs
|
||||
<Tabs>
|
||||
<Tab title="Honcho API">
|
||||
<Update label="v3.0.8 (Current)">
|
||||
<Update label="v3.0.9 (Current)">
|
||||
### Changed
|
||||
|
||||
- Connection acquisition is now a single attempt with no server-side retry, on a vanilla `AsyncSession`. A new `DB_CONNECT_TIMEOUT_SECONDS` (default 2s) bounds the attempt so a saturated or unreachable pooler fails fast instead of holding a client connection open to re-knock. A saturated DB now surfaces to the caller — the API returns an error and the deriver backs off and retries on a later poll — which lets the pooler drain rather than amplifying saturation.
|
||||
|
||||
### Added
|
||||
|
||||
- Deriver poll jitter so instances that start together don't poll in lockstep: `DERIVER_POLLING_STARTUP_JITTER_SECONDS` (random delay before the first poll, default 30s) and `DERIVER_POLLING_JITTER_RATIO` (±fraction applied to every poll sleep, default 0.5). Both disable at `0.0`; the underlying backoff schedule is unchanged.
|
||||
|
||||
### Removed
|
||||
|
||||
- Reverted the connection-checkout retry and `HonchoAsyncSession` custom session introduced in 3.0.8. Removed the `DB_CONNECTION_RETRY_ENABLED` / `DB_CONNECTION_RETRY_MAX_DELAY_SECONDS` / `DB_CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS` / `DB_CONNECTION_RETRY_BACKOFF_MAX_SECONDS` settings, the `db_connection_acquisitions{outcome=...}` Prometheus counter, and the `db.pool.acquire` Sentry span. Alerting built on `db_connection_acquisitions` should migrate to `db_pool_connections` / `db_queries_in_flight`.
|
||||
</Update>
|
||||
|
||||
<Update label="v3.0.8">
|
||||
### Added
|
||||
|
||||
- Connection-checkout retry with bounded exponential backoff (tenacity) on `get_db`/`tracked_db`: transient transaction-pooler (Supavisor) rejections — SQLAlchemy `TimeoutError` and `OperationalError` — now retry with backoff instead of surfacing as 500s under client-connection saturation. Gated by
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@
|
|||
"navigation": {
|
||||
"versions": [
|
||||
{
|
||||
"version": "v3.0.8",
|
||||
"version": "v3.0.9",
|
||||
"api": {
|
||||
"openapi": ["v3/openapi.json"]
|
||||
},
|
||||
|
|
|
|||
|
|
@ -360,6 +360,11 @@ DERIVER_POLLING_SLEEP_INTERVAL_SECONDS=1.0
|
|||
DERIVER_POLLING_BACKOFF_ENABLED=true
|
||||
DERIVER_POLLING_SLEEP_MAX_INTERVAL_SECONDS=30.0
|
||||
DERIVER_POLLING_BACKOFF_MULTIPLIER=2.0
|
||||
# Jitter so instances that start together don't poll in lockstep. Startup: sleep
|
||||
# a random delay in [0, value] before the first poll (0.0 disables). Per-cycle:
|
||||
# multiply every poll sleep by a random factor in [1-ratio, 1+ratio] (0.0 disables).
|
||||
DERIVER_POLLING_STARTUP_JITTER_SECONDS=30.0
|
||||
DERIVER_POLLING_JITTER_RATIO=0.5
|
||||
DERIVER_STALE_SESSION_TIMEOUT_MINUTES=5
|
||||
|
||||
# Queue management
|
||||
|
|
@ -469,12 +474,9 @@ DB_POOL_TIMEOUT=5
|
|||
DB_POOL_RECYCLE=300
|
||||
DB_POOL_USE_LIFO=true
|
||||
DB_SQL_DEBUG=false
|
||||
# Bounded retry around connection checkout. For a pooled (non-null) DB_POOL_CLASS,
|
||||
# DB_POOL_TIMEOUT must be < DB_CONNECTION_RETRY_MAX_DELAY_SECONDS or startup fails.
|
||||
DB_CONNECTION_RETRY_ENABLED=true
|
||||
DB_CONNECTION_RETRY_MAX_DELAY_SECONDS=10.0
|
||||
DB_CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS=0.1
|
||||
DB_CONNECTION_RETRY_BACKOFF_MAX_SECONDS=2.0
|
||||
# Per-connection establish timeout (seconds) so a single connection attempt
|
||||
# fails fast instead of hanging when the server/pooler is unreachable.
|
||||
DB_CONNECT_TIMEOUT_SECONDS=2
|
||||
```
|
||||
|
||||
### Authentication
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "honcho"
|
||||
version = "3.0.8"
|
||||
version = "3.0.9"
|
||||
description = "Honcho Server"
|
||||
authors = [
|
||||
{name = "Plastic Labs", email = "hello@plasticlabs.ai"},
|
||||
|
|
|
|||
|
|
@ -613,8 +613,8 @@ class DBSettings(HonchoSettings):
|
|||
POOL_SIZE: Annotated[int, Field(default=10, gt=0, le=1000)] = 10
|
||||
MAX_OVERFLOW: Annotated[int, Field(default=20, ge=0, le=1000)] = 20
|
||||
POOL_TIMEOUT: Annotated[int, Field(default=5, gt=0, le=300)] = (
|
||||
5 # seconds; kept under CONNECTION_RETRY_MAX_DELAY_SECONDS so a pool
|
||||
# checkout fails fast enough to allow a retry within the budget
|
||||
5 # seconds a pooled checkout may wait for a free connection (QueuePool
|
||||
# only; NullPool has no local queue wait)
|
||||
)
|
||||
POOL_RECYCLE: Annotated[int, Field(default=300, gt=0, le=7200)] = (
|
||||
300 # seconds (max 2 hours)
|
||||
|
|
@ -623,47 +623,12 @@ class DBSettings(HonchoSettings):
|
|||
SQL_DEBUG: bool = False
|
||||
TRACING: bool = False
|
||||
|
||||
# Bounded exponential-backoff retry around connection acquisition. Applied
|
||||
# lazily on the first DB use of any session (HonchoAsyncSession) — both the
|
||||
# request path and background/tracked_db scopes — without forcing an eager
|
||||
# checkout. Guards against transient transaction-pooler saturation (e.g.
|
||||
# Supavisor rejecting with "too many clients") by retrying the checkout
|
||||
# instead of failing immediately. CONNECTION_RETRY_MAX_DELAY_SECONDS is the
|
||||
# TOTAL retry
|
||||
# budget; with a real QueuePool, a single checkout can block up to
|
||||
# POOL_TIMEOUT, so POOL_TIMEOUT must stay below the budget for a retry to be
|
||||
# possible (enforced below). With NullPool (the transaction-pooler setup)
|
||||
# there is no local queue wait, so saturation is a fast OperationalError and
|
||||
# POOL_TIMEOUT does not apply.
|
||||
CONNECTION_RETRY_ENABLED: bool = True
|
||||
CONNECTION_RETRY_MAX_DELAY_SECONDS: Annotated[
|
||||
float, Field(default=10.0, gt=0.0, le=120.0)
|
||||
] = 10.0
|
||||
CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS: Annotated[
|
||||
float, Field(default=0.1, gt=0.0, le=10.0)
|
||||
] = 0.1
|
||||
CONNECTION_RETRY_BACKOFF_MAX_SECONDS: Annotated[
|
||||
float, Field(default=2.0, gt=0.0, le=30.0)
|
||||
] = 2.0
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_retry_budget_vs_pool_timeout(self) -> "DBSettings":
|
||||
# Only meaningful for a real local pool: NullPool has no queue wait, so
|
||||
# POOL_TIMEOUT is unused there. For a QueuePool, a checkout can block up
|
||||
# to POOL_TIMEOUT, so it must be < the total retry budget or the first
|
||||
# attempt consumes the whole budget and no retry ever happens.
|
||||
if (
|
||||
self.POOL_CLASS != "null"
|
||||
and self.CONNECTION_RETRY_ENABLED
|
||||
and self.POOL_TIMEOUT >= self.CONNECTION_RETRY_MAX_DELAY_SECONDS
|
||||
):
|
||||
raise ValueError(
|
||||
f"DB_POOL_TIMEOUT ({self.POOL_TIMEOUT}s) must be less than "
|
||||
+ "DB_CONNECTION_RETRY_MAX_DELAY_SECONDS "
|
||||
+ f"({self.CONNECTION_RETRY_MAX_DELAY_SECONDS}s) so a pooled checkout "
|
||||
+ "can fail fast enough to be retried within the budget."
|
||||
)
|
||||
return self
|
||||
# Per-connection establish timeout (seconds) passed to the driver, so a
|
||||
# single connection attempt fails fast instead of hanging when the server or
|
||||
# pooler is unreachable or stalled. Connection acquisition is a single
|
||||
# attempt with no retry; callers handle failure (the API surfaces it, the
|
||||
# deriver backs off and retries on a later poll).
|
||||
CONNECT_TIMEOUT_SECONDS: Annotated[int, Field(default=2, gt=0, le=60)] = 2
|
||||
|
||||
|
||||
class AuthSettings(HonchoSettings):
|
||||
|
|
@ -792,6 +757,17 @@ class DeriverSettings(HonchoSettings):
|
|||
POLLING_BACKOFF_MULTIPLIER: Annotated[
|
||||
float, Field(default=2.0, ge=1.0, le=10.0)
|
||||
] = 2.0
|
||||
# Sleep a uniform-random delay in [0, POLLING_STARTUP_JITTER_SECONDS] before
|
||||
# the first poll so instances that start together don't poll in lockstep.
|
||||
# Set to 0.0 to disable.
|
||||
POLLING_STARTUP_JITTER_SECONDS: Annotated[
|
||||
float, Field(default=30.0, ge=0.0, le=300.0)
|
||||
] = 30.0
|
||||
# Multiply every poll sleep by a random factor in [1 - ratio, 1 + ratio]
|
||||
# (0.5 -> [0.5x, 1.5x]) so poll loops don't re-converge over time. The
|
||||
# backoff schedule is unchanged; only the returned sleep is scattered. Set
|
||||
# to 0.0 to disable.
|
||||
POLLING_JITTER_RATIO: Annotated[float, Field(default=0.5, ge=0.0, le=1.0)] = 0.5
|
||||
STALE_SESSION_TIMEOUT_MINUTES: Annotated[int, Field(default=5, gt=0, le=1440)] = 5
|
||||
|
||||
# Retention window (seconds) for keeping errored items in the queue
|
||||
|
|
|
|||
257
src/db.py
257
src/db.py
|
|
@ -2,10 +2,7 @@ import contextvars
|
|||
import logging
|
||||
from typing import Any
|
||||
|
||||
import sentry_sdk
|
||||
from sqlalchemy import MetaData, event, text
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlalchemy.exc import TimeoutError as SQLAlchemyTimeoutError
|
||||
from sqlalchemy.ext.asyncio import (
|
||||
AsyncSession,
|
||||
async_sessionmaker,
|
||||
|
|
@ -13,22 +10,18 @@ from sqlalchemy.ext.asyncio import (
|
|||
)
|
||||
from sqlalchemy.orm import declarative_base
|
||||
from sqlalchemy.pool import NullPool, QueuePool
|
||||
from tenacity import (
|
||||
AsyncRetrying,
|
||||
retry_if_exception_type,
|
||||
stop_after_delay,
|
||||
wait_exponential_jitter,
|
||||
)
|
||||
|
||||
from src.config import settings
|
||||
from src.telemetry.prometheus.metrics import (
|
||||
db_queries_in_flight_gauge,
|
||||
prometheus_metrics,
|
||||
)
|
||||
from src.telemetry.prometheus.metrics import db_queries_in_flight_gauge
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
connect_args = {"prepare_threshold": None}
|
||||
connect_args = {
|
||||
"prepare_threshold": None,
|
||||
# Bound a single connection attempt so it fails fast instead of hanging when
|
||||
# the server/pooler is unreachable or stalled (psycopg, seconds).
|
||||
"connect_timeout": settings.DB.CONNECT_TIMEOUT_SECONDS,
|
||||
}
|
||||
|
||||
# Context variable to store request context
|
||||
request_context: contextvars.ContextVar[str | None] = contextvars.ContextVar(
|
||||
|
|
@ -59,25 +52,46 @@ engine = create_async_engine(
|
|||
**engine_kwargs,
|
||||
)
|
||||
|
||||
# NOTE: SessionLocal is defined further down, after HonchoAsyncSession (its
|
||||
# session class) and acquire_connection_with_retry (which that class calls).
|
||||
|
||||
# Errors worth retrying when acquiring a pooled connection: SQLAlchemy's local
|
||||
# pool-checkout timeout, and OperationalError (how a saturated transaction
|
||||
# pooler surfaces "too many clients" / connection refusals).
|
||||
RETRYABLE_DB_CONNECTION_ERRORS = (SQLAlchemyTimeoutError, OperationalError)
|
||||
|
||||
# Identifies this process ("api" | "deriver") on DB metrics. Set once at startup
|
||||
# by register_db_query_instrumentation; stays "unknown" if metrics are disabled.
|
||||
_db_instance_type: str = "unknown"
|
||||
# A vanilla AsyncSession is lazy: it checks out a pooled connection on the first
|
||||
# DB-touching call (not at construction) and couples the checkout to the
|
||||
# statement, so a handler doing non-DB work (embedding/file/LLM) before its
|
||||
# first query does not pin a connection across it. Connection acquisition is a
|
||||
# single attempt with no retry — callers handle a saturated/unreachable DB (the
|
||||
# API surfaces the error; the deriver backs off and retries on a later poll).
|
||||
SessionLocal = async_sessionmaker(
|
||||
autocommit=False,
|
||||
autoflush=False,
|
||||
expire_on_commit=False,
|
||||
bind=engine,
|
||||
class_=AsyncSession,
|
||||
)
|
||||
|
||||
|
||||
def _record_acquisition_outcome(outcome: str) -> None:
|
||||
"""Record a connection-acquisition outcome (no-op when metrics disabled)."""
|
||||
if settings.METRICS.ENABLED:
|
||||
prometheus_metrics.record_db_connection_acquisition(
|
||||
instance_type=_db_instance_type, outcome=outcome
|
||||
)
|
||||
def _set_application_name_on_checkout(
|
||||
dbapi_connection: Any, _connection_record: Any, _connection_proxy: Any
|
||||
) -> None:
|
||||
"""Tag each checked-out connection with the current request context.
|
||||
|
||||
Registered only when ``DB.TRACING`` is on. Fires on every pool checkout (so a
|
||||
reused pooled connection is re-tagged for the new caller), reading the
|
||||
per-task ``request_context`` the request/task scope has already set.
|
||||
Best-effort: a failure here must never break the checkout.
|
||||
"""
|
||||
context = request_context.get() or "unknown"
|
||||
try:
|
||||
cursor = dbapi_connection.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
"SELECT set_config('application_name', %s, false)", (context,)
|
||||
)
|
||||
finally:
|
||||
cursor.close()
|
||||
except Exception:
|
||||
logger.debug("setting application_name on checkout failed", exc_info=True)
|
||||
|
||||
|
||||
if settings.DB.TRACING:
|
||||
event.listen(engine.sync_engine, "checkout", _set_application_name_on_checkout)
|
||||
|
||||
|
||||
def get_pool_stats() -> dict[str, int]:
|
||||
|
|
@ -107,184 +121,6 @@ def get_pool_stats() -> dict[str, int]:
|
|||
return zeros
|
||||
|
||||
|
||||
async def acquire_connection_with_retry(db: AsyncSession, context: str) -> None:
|
||||
"""Force pool checkout (which ``SessionLocal()`` defers) with bounded backoff.
|
||||
|
||||
``SessionLocal()`` is lazy: the pool checkout — and any pooler rejection —
|
||||
happens on the first query. We force it here inside a retry block so that
|
||||
transient saturation of the transaction pooler is retried with exponential
|
||||
backoff rather than surfacing as an immediate error. The checkout is wrapped
|
||||
in a Sentry span so wait time is visible in traces; on budget exhaustion the
|
||||
original error is reraised after capturing live pool stats to Sentry.
|
||||
|
||||
Each attempt rolls the session back on a retryable failure before retrying:
|
||||
a failed checkout can leave the autobegun transaction in a pending-rollback
|
||||
state, which would make the next ``db.connection()`` raise instead of
|
||||
re-checking-out cleanly. The rollback is pure Python-side state cleanup when
|
||||
no connection was bound, so it is cheap and safe.
|
||||
"""
|
||||
with sentry_sdk.start_span(op="db.pool.acquire", name=context):
|
||||
if not settings.DB.CONNECTION_RETRY_ENABLED:
|
||||
await db.connection()
|
||||
return
|
||||
attempts = 0
|
||||
try:
|
||||
async for attempt in AsyncRetrying(
|
||||
wait=wait_exponential_jitter(
|
||||
initial=settings.DB.CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS,
|
||||
max=settings.DB.CONNECTION_RETRY_BACKOFF_MAX_SECONDS,
|
||||
),
|
||||
stop=stop_after_delay(settings.DB.CONNECTION_RETRY_MAX_DELAY_SECONDS),
|
||||
retry=retry_if_exception_type(RETRYABLE_DB_CONNECTION_ERRORS),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
attempts += 1
|
||||
try:
|
||||
await db.connection()
|
||||
except RETRYABLE_DB_CONNECTION_ERRORS:
|
||||
# Reset session state so the next attempt starts clean.
|
||||
try:
|
||||
await db.rollback()
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"rollback after failed checkout failed",
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
except RETRYABLE_DB_CONNECTION_ERRORS as e:
|
||||
_record_acquisition_outcome("exhausted")
|
||||
if settings.SENTRY.ENABLED:
|
||||
sentry_sdk.set_context("db_pool", get_pool_stats())
|
||||
sentry_sdk.capture_exception(e)
|
||||
raise
|
||||
# "ok" on first try, "retried" if backoff was needed before success.
|
||||
_record_acquisition_outcome("ok" if attempts <= 1 else "retried")
|
||||
|
||||
|
||||
class HonchoAsyncSession(AsyncSession):
|
||||
"""AsyncSession that lazily checks out its connection, with retry.
|
||||
|
||||
The pool checkout — and any pooler-rejection retry via
|
||||
``acquire_connection_with_retry`` — happens on the FIRST DB-touching call,
|
||||
not at construction. So a request handler that does non-DB work (embedding,
|
||||
file processing, an LLM call) before its first query does NOT pin a
|
||||
connection across that work, while still getting checkout retry on the
|
||||
request path.
|
||||
|
||||
Only the checkout is retried; the SQL statement itself runs exactly once
|
||||
(we never retry ``super().execute`` after a broad OperationalError), so
|
||||
writes are never duplicated. The context for tracing is read from the
|
||||
``request_context`` ContextVar, which the request/task scope has already set.
|
||||
"""
|
||||
|
||||
# Class-level default; per-instance assignment shadows it (the subclass has
|
||||
# a __dict__ even though AsyncSession declares __slots__).
|
||||
_honcho_acquired: bool = False
|
||||
|
||||
async def _ensure_acquired(self) -> None:
|
||||
if self._honcho_acquired:
|
||||
return
|
||||
context = request_context.get() or "unknown"
|
||||
await acquire_connection_with_retry(self, context)
|
||||
self._honcho_acquired = True
|
||||
if settings.DB.TRACING:
|
||||
# Forced checkout already happened above; this rides the same
|
||||
# connection. super() to avoid re-entering _ensure_acquired.
|
||||
await super().execute(
|
||||
text("SELECT set_config('application_name', :name, false)"),
|
||||
{"name": context},
|
||||
)
|
||||
|
||||
# The overrides below are thin: ensure the connection is checked out (once,
|
||||
# with retry) before delegating to AsyncSession. They cover every public
|
||||
# DB-touching async method so the "lazy retry on first DB use" guarantee has
|
||||
# no holes. Signatures are widened to *args/**kwargs because we only forward;
|
||||
# call sites are typed against the AsyncSession base, so this does not weaken
|
||||
# type-checking elsewhere. (connection() is intentionally NOT wrapped —
|
||||
# acquire_connection_with_retry calls it, so wrapping would recurse.)
|
||||
async def execute(self, *args: Any, **kwargs: Any) -> Any:
|
||||
await self._ensure_acquired()
|
||||
return await super().execute(*args, **kwargs)
|
||||
|
||||
async def scalar(self, *args: Any, **kwargs: Any) -> Any:
|
||||
await self._ensure_acquired()
|
||||
return await super().scalar(*args, **kwargs)
|
||||
|
||||
async def scalars(self, *args: Any, **kwargs: Any) -> Any:
|
||||
await self._ensure_acquired()
|
||||
return await super().scalars(*args, **kwargs)
|
||||
|
||||
async def get(self, *args: Any, **kwargs: Any) -> Any:
|
||||
await self._ensure_acquired()
|
||||
return await super().get(*args, **kwargs)
|
||||
|
||||
async def get_one(self, *args: Any, **kwargs: Any) -> Any:
|
||||
await self._ensure_acquired()
|
||||
return await super().get_one(*args, **kwargs)
|
||||
|
||||
async def stream(self, *args: Any, **kwargs: Any) -> Any:
|
||||
await self._ensure_acquired()
|
||||
return await super().stream(*args, **kwargs)
|
||||
|
||||
async def stream_scalars(self, *args: Any, **kwargs: Any) -> Any:
|
||||
await self._ensure_acquired()
|
||||
return await super().stream_scalars(*args, **kwargs)
|
||||
|
||||
async def flush(self, *args: Any, **kwargs: Any) -> None:
|
||||
await self._ensure_acquired()
|
||||
await super().flush(*args, **kwargs)
|
||||
|
||||
async def merge(self, *args: Any, **kwargs: Any) -> Any:
|
||||
await self._ensure_acquired()
|
||||
return await super().merge(*args, **kwargs)
|
||||
|
||||
async def refresh(self, *args: Any, **kwargs: Any) -> None:
|
||||
await self._ensure_acquired()
|
||||
await super().refresh(*args, **kwargs)
|
||||
|
||||
async def delete(self, *args: Any, **kwargs: Any) -> None:
|
||||
await self._ensure_acquired()
|
||||
await super().delete(*args, **kwargs)
|
||||
|
||||
async def commit(self) -> None:
|
||||
# Ensures the add()->commit() path (autoflush on commit) also retries.
|
||||
await self._ensure_acquired()
|
||||
try:
|
||||
await super().commit()
|
||||
finally:
|
||||
# Transaction ended; a later op must re-acquire (and re-wrap retry).
|
||||
self._honcho_acquired = False
|
||||
|
||||
async def rollback(self) -> None:
|
||||
try:
|
||||
await super().rollback()
|
||||
finally:
|
||||
self._honcho_acquired = False
|
||||
|
||||
async def close(self) -> None:
|
||||
try:
|
||||
await super().close()
|
||||
finally:
|
||||
# The connection is released; a reused session must re-acquire.
|
||||
self._honcho_acquired = False
|
||||
|
||||
async def reset(self) -> None:
|
||||
try:
|
||||
await super().reset()
|
||||
finally:
|
||||
self._honcho_acquired = False
|
||||
|
||||
|
||||
SessionLocal = async_sessionmaker(
|
||||
autocommit=False,
|
||||
autoflush=False,
|
||||
expire_on_commit=False,
|
||||
bind=engine,
|
||||
class_=HonchoAsyncSession,
|
||||
)
|
||||
|
||||
|
||||
class DBQueryInflightTracker:
|
||||
"""Tracks statements executing on the wire via SQLAlchemy cursor events.
|
||||
|
||||
|
|
@ -340,8 +176,7 @@ def register_db_query_instrumentation(instance_type: str) -> None:
|
|||
re-run lifespan or test startup) won't attach duplicate listeners, which
|
||||
would double-count in-flight statements.
|
||||
"""
|
||||
global _db_instance_type, _inflight_tracker, _db_query_instrumentation_registered
|
||||
_db_instance_type = instance_type
|
||||
global _inflight_tracker, _db_query_instrumentation_registered
|
||||
if not settings.METRICS.ENABLED or _db_query_instrumentation_registered:
|
||||
return
|
||||
child = db_queries_in_flight_gauge.labels(instance_type=instance_type)
|
||||
|
|
|
|||
|
|
@ -10,11 +10,10 @@ from src.db import SessionLocal, request_context
|
|||
async def get_db():
|
||||
"""FastAPI Dependency Generator for Database.
|
||||
|
||||
The session is lazy: it does NOT check out a pooled connection here.
|
||||
HonchoAsyncSession acquires it (with retry) on the first DB-touching call,
|
||||
so a handler doing non-DB work (embedding/file/LLM) before its first query
|
||||
does not pin a connection across it. Tracing's application_name and the
|
||||
checkout retry both live in the session, keyed off request_context.
|
||||
The session is lazy: it does NOT check out a pooled connection here. The
|
||||
AsyncSession checks one out on the first DB-touching call, so a handler doing
|
||||
non-DB work (embedding/file/LLM) before its first query does not pin a
|
||||
connection across it.
|
||||
"""
|
||||
db: AsyncSession = SessionLocal()
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import random
|
||||
import signal
|
||||
from asyncio import Task
|
||||
from collections.abc import Sequence
|
||||
|
|
@ -205,6 +207,7 @@ class QueueManager:
|
|||
# Run the polling loop directly in this task
|
||||
logger.debug("Starting polling loop directly")
|
||||
try:
|
||||
await self._sleep_startup_jitter()
|
||||
await self.polling_loop()
|
||||
finally:
|
||||
await self.cleanup()
|
||||
|
|
@ -391,6 +394,35 @@ class QueueManager:
|
|||
"""Snap the polling interval back to the base after finding work."""
|
||||
self._current_poll_interval = settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS
|
||||
|
||||
def _jitter(self, seconds: float) -> float:
|
||||
"""Scatter a sleep by +/- POLLING_JITTER_RATIO to avoid lockstep polling.
|
||||
|
||||
Returns a uniform-random value in [(1-ratio)*seconds, (1+ratio)*seconds].
|
||||
Only the returned sleep is scattered; the underlying backoff schedule is
|
||||
left unchanged. A ratio of 0.0 returns ``seconds`` unchanged.
|
||||
"""
|
||||
ratio = settings.DERIVER.POLLING_JITTER_RATIO
|
||||
if ratio <= 0.0:
|
||||
return seconds
|
||||
# Scheduling jitter, not security/crypto — stdlib random is appropriate.
|
||||
return seconds * random.uniform(1.0 - ratio, 1.0 + ratio) # nosec B311
|
||||
|
||||
async def _sleep_startup_jitter(self) -> None:
|
||||
"""Sleep a random delay before the first poll so instances that start
|
||||
together don't poll in lockstep. Interruptible by shutdown so a signal
|
||||
during the delay exits promptly. No-op when the window is 0.0.
|
||||
"""
|
||||
window = settings.DERIVER.POLLING_STARTUP_JITTER_SECONDS
|
||||
if window <= 0.0:
|
||||
return
|
||||
# Scheduling jitter, not security/crypto — stdlib random is appropriate.
|
||||
delay = random.uniform(0.0, window) # nosec B311
|
||||
logger.debug(f"Startup poll jitter: sleeping {delay:.1f}s before first poll")
|
||||
# Timeout (slept the full delay without a shutdown) is the normal path;
|
||||
# an early return means shutdown fired and polling_loop will exit at once.
|
||||
with contextlib.suppress(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(self.shutdown_event.wait(), timeout=delay)
|
||||
|
||||
def _advance_poll_interval(self) -> float:
|
||||
"""Return the current idle/backoff sleep, then grow it toward the cap."""
|
||||
interval = self._current_poll_interval
|
||||
|
|
@ -400,7 +432,7 @@ class QueueManager:
|
|||
* settings.DERIVER.POLLING_BACKOFF_MULTIPLIER,
|
||||
settings.DERIVER.POLLING_SLEEP_MAX_INTERVAL_SECONDS,
|
||||
)
|
||||
return interval
|
||||
return self._jitter(interval)
|
||||
|
||||
async def polling_loop(self) -> None:
|
||||
"""Main polling loop to find and process new work units"""
|
||||
|
|
@ -419,7 +451,9 @@ class QueueManager:
|
|||
# when capacity frees rather than backing off.
|
||||
if self.semaphore.locked():
|
||||
# logger.debug("All workers busy, waiting")
|
||||
await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS)
|
||||
await asyncio.sleep(
|
||||
self._jitter(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS)
|
||||
)
|
||||
continue
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -119,8 +119,7 @@ if SENTRY_ENABLED:
|
|||
FastApiIntegration(
|
||||
transaction_style="endpoint",
|
||||
),
|
||||
# Explicit so DB-query spans (and our db.pool.acquire span) are not
|
||||
# reliant on auto-enabling.
|
||||
# Explicit so DB-query spans are not reliant on auto-enabling.
|
||||
SqlalchemyIntegration(),
|
||||
],
|
||||
before_send=before_send,
|
||||
|
|
|
|||
|
|
@ -127,17 +127,9 @@ telemetry_buffer_size_gauge = NamespacedGauge(
|
|||
["namespace"],
|
||||
)
|
||||
|
||||
# DB connection-pool health. The acquisitions counter measures how often we hit
|
||||
# (and retry through) transaction-pooler saturation; the in-flight gauge counts
|
||||
# statements actually executing on the wire, so checked_out minus in_flight
|
||||
# reveals connections held but parked (the "idle in transaction during an
|
||||
# external call" antipattern).
|
||||
db_connection_acquisitions_counter = NamespacedCounter(
|
||||
"db_connection_acquisitions",
|
||||
"DB connection acquisitions by outcome (ok=first try, retried, exhausted)",
|
||||
["namespace", "instance_type", "outcome"],
|
||||
)
|
||||
|
||||
# DB connection-pool health. The in-flight gauge counts statements actually
|
||||
# executing on the wire, so checked_out minus in_flight reveals connections held
|
||||
# but parked (the "idle in transaction during an external call" antipattern).
|
||||
db_queries_in_flight_gauge = NamespacedGauge(
|
||||
"db_queries_in_flight",
|
||||
"DB statements currently executing on a connection for this instance",
|
||||
|
|
@ -294,18 +286,6 @@ class PrometheusMetrics:
|
|||
except Exception as e:
|
||||
self._handle_metric_error("set_telemetry_buffer_size", e)
|
||||
|
||||
def record_db_connection_acquisition(
|
||||
self, *, instance_type: str, outcome: str
|
||||
) -> None:
|
||||
# outcome is one of "ok" | "retried" | "exhausted".
|
||||
try:
|
||||
db_connection_acquisitions_counter.labels(
|
||||
instance_type=instance_type,
|
||||
outcome=outcome,
|
||||
).inc()
|
||||
except Exception as e:
|
||||
self._handle_metric_error("record_db_connection_acquisition", e)
|
||||
|
||||
|
||||
prometheus_metrics = PrometheusMetrics()
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
|
@ -1583,3 +1584,62 @@ class TestQueueProcessing:
|
|||
claimed = await qm.get_and_claim_work_units()
|
||||
assert work_unit_key is not None
|
||||
assert work_unit_key in claimed
|
||||
|
||||
|
||||
class TestPollingJitter:
|
||||
"""Polling jitter: desynchronize poll loops without changing the schedule."""
|
||||
|
||||
def test_jitter_stays_within_ratio_bounds(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.5)
|
||||
qm = QueueManager()
|
||||
samples = [qm._jitter(10.0) for _ in range(1000)] # pyright: ignore[reportPrivateUsage]
|
||||
assert all(5.0 <= s <= 15.0 for s in samples)
|
||||
# A 0.5 ratio over 1000 samples should produce real spread, not a constant.
|
||||
assert max(samples) - min(samples) > 1.0
|
||||
|
||||
def test_jitter_ratio_zero_is_identity(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0)
|
||||
qm = QueueManager()
|
||||
assert all(qm._jitter(7.5) == 7.5 for _ in range(50)) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
def test_advance_jitters_return_but_keeps_deterministic_schedule(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.5)
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", True)
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0)
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0)
|
||||
monkeypatch.setattr(
|
||||
settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 30.0
|
||||
)
|
||||
qm = QueueManager()
|
||||
|
||||
# The underlying schedule advances 1 -> 2 -> 4 -> ... -> 30 deterministically;
|
||||
# each returned sleep is jittered within [0.5x, 1.5x] of the pre-advance step.
|
||||
expected_schedule = [1.0, 2.0, 4.0, 8.0, 16.0, 30.0, 30.0]
|
||||
for step in expected_schedule:
|
||||
returned = qm._advance_poll_interval() # pyright: ignore[reportPrivateUsage]
|
||||
assert 0.5 * step <= returned <= 1.5 * step
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_startup_jitter_disabled_returns_immediately(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_STARTUP_JITTER_SECONDS", 0.0)
|
||||
qm = QueueManager()
|
||||
# Window 0.0 must not sleep at all.
|
||||
await qm._sleep_startup_jitter() # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_startup_jitter_interrupted_by_shutdown(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_STARTUP_JITTER_SECONDS", 300.0)
|
||||
qm = QueueManager()
|
||||
qm.shutdown_event.set()
|
||||
# A shutdown already signalled must short-circuit the (long) jitter sleep.
|
||||
await asyncio.wait_for(qm._sleep_startup_jitter(), timeout=1.0) # pyright: ignore[reportPrivateUsage]
|
||||
|
|
|
|||
|
|
@ -1,248 +1,96 @@
|
|||
"""Unit tests for DB connection resilience + observability.
|
||||
|
||||
These are DB-free: they exercise the retry helper against a fake session, the
|
||||
deriver polling backoff math, and the in-flight gauge listeners directly.
|
||||
These are DB-free: they exercise the application_name checkout hook against a
|
||||
fake DBAPI connection, the deriver polling backoff math, and the in-flight gauge
|
||||
listeners directly.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
import src.db as db_module
|
||||
from src.config import settings
|
||||
from src.db import DBQueryInflightTracker, acquire_connection_with_retry
|
||||
from src.telemetry.prometheus.metrics import (
|
||||
db_connection_acquisitions_counter,
|
||||
db_queries_in_flight_gauge,
|
||||
)
|
||||
from src.db import DBQueryInflightTracker
|
||||
from src.telemetry.prometheus.metrics import db_queries_in_flight_gauge
|
||||
|
||||
|
||||
def _make_operational_error() -> OperationalError:
|
||||
"""A stand-in for how a saturated pooler surfaces ('too many clients')."""
|
||||
return OperationalError("SELECT 1", {}, Exception("too many clients"))
|
||||
|
||||
|
||||
class _FlakyConnSession:
|
||||
"""Fake AsyncSession whose connection() fails N times then succeeds."""
|
||||
|
||||
def __init__(self, fail_times: int, *, always_fail: bool = False) -> None:
|
||||
self.fail_times: int = fail_times
|
||||
self.always_fail: bool = always_fail
|
||||
self.calls: int = 0
|
||||
self.rollback_calls: int = 0
|
||||
|
||||
async def connection(self) -> None:
|
||||
self.calls += 1
|
||||
if self.always_fail or self.calls <= self.fail_times:
|
||||
raise _make_operational_error()
|
||||
|
||||
async def rollback(self) -> None:
|
||||
self.rollback_calls += 1
|
||||
|
||||
|
||||
def _acq_count(outcome: str) -> float:
|
||||
child = db_connection_acquisitions_counter.labels(
|
||||
instance_type="api", outcome=outcome
|
||||
)
|
||||
return float(child._value.get()) # pyright: ignore[reportPrivateUsage, reportUnknownArgumentType]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def metrics_on(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(settings.METRICS, "ENABLED", True)
|
||||
monkeypatch.setattr(settings.METRICS, "NAMESPACE", "test")
|
||||
monkeypatch.setattr(db_module, "_db_instance_type", "api")
|
||||
# Fast, deterministic backoff so retry tests don't sleep.
|
||||
monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_ENABLED", True)
|
||||
monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS", 0.001)
|
||||
monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_BACKOFF_MAX_SECONDS", 0.01)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("metrics_on")
|
||||
async def test_acquire_succeeds_first_try_records_ok() -> None:
|
||||
before = _acq_count("ok")
|
||||
session = _FlakyConnSession(fail_times=0)
|
||||
await acquire_connection_with_retry(session, "request:test") # pyright: ignore[reportArgumentType]
|
||||
assert session.calls == 1
|
||||
assert _acq_count("ok") == before + 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("metrics_on")
|
||||
async def test_acquire_retries_then_succeeds_records_retried() -> None:
|
||||
before = _acq_count("retried")
|
||||
session = _FlakyConnSession(fail_times=2)
|
||||
await acquire_connection_with_retry(session, "request:test") # pyright: ignore[reportArgumentType]
|
||||
assert session.calls == 3 # two failures then success
|
||||
# Session is reset after each failed checkout before the next attempt.
|
||||
assert session.rollback_calls == 2
|
||||
assert _acq_count("retried") == before + 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("metrics_on")
|
||||
async def test_acquire_exhausts_budget_reraises_and_records(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# Tiny budget so the loop gives up quickly under sustained failure.
|
||||
monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_MAX_DELAY_SECONDS", 0.05)
|
||||
before = _acq_count("exhausted")
|
||||
session = _FlakyConnSession(fail_times=0, always_fail=True)
|
||||
with pytest.raises(OperationalError):
|
||||
await acquire_connection_with_retry(session, "request:test") # pyright: ignore[reportArgumentType]
|
||||
assert session.calls >= 1
|
||||
assert _acq_count("exhausted") == before + 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_disabled_calls_once(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_ENABLED", False)
|
||||
session = _FlakyConnSession(fail_times=0)
|
||||
await acquire_connection_with_retry(session, "request:test") # pyright: ignore[reportArgumentType]
|
||||
assert session.calls == 1
|
||||
|
||||
|
||||
# --- HonchoAsyncSession: lazy checkout with retry on first DB use ------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_lazy_acquires_on_first_db_use(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""No checkout at construction; acquisition (with retry) happens once on the
|
||||
first execute, and the statement itself runs exactly once."""
|
||||
monkeypatch.setattr(settings.DB, "TRACING", False)
|
||||
acquired: list[str] = []
|
||||
executed: list[Any] = []
|
||||
|
||||
async def fake_acquire(_session: Any, context: str) -> None:
|
||||
acquired.append(context)
|
||||
|
||||
async def fake_super_execute(_self: Any, *args: Any, **_kw: Any) -> str:
|
||||
executed.append(args[0] if args else None)
|
||||
return "result"
|
||||
|
||||
monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire)
|
||||
monkeypatch.setattr(AsyncSession, "execute", fake_super_execute)
|
||||
def test_session_local_uses_vanilla_async_session() -> None:
|
||||
"""Regression guard: no custom session subclass / acquisition logic.
|
||||
|
||||
Connection acquisition is a single lazy checkout owned by AsyncSession; there
|
||||
must be no re-introduced eager-checkout or retry hooks on the session.
|
||||
"""
|
||||
session = db_module.SessionLocal()
|
||||
assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage]
|
||||
assert acquired == [] # constructing the session does NOT check out
|
||||
|
||||
result = await session.execute("SELECT 1")
|
||||
assert result == "result"
|
||||
assert acquired == ["unknown"] # acquired exactly once, on first use
|
||||
assert session._honcho_acquired is True # pyright: ignore[reportPrivateUsage]
|
||||
assert executed == ["SELECT 1"] # statement ran once (not retried)
|
||||
|
||||
await session.execute("SELECT 2")
|
||||
assert acquired == ["unknown"] # still once — no re-acquire on later use
|
||||
assert type(session) is AsyncSession
|
||||
assert not hasattr(session, "_ensure_acquired")
|
||||
assert not hasattr(session, "_honcho_acquired")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_tracing_sets_application_name_on_acquire(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(settings.DB, "TRACING", True)
|
||||
calls: list[tuple[Any, ...]] = []
|
||||
# --- application_name checkout hook ------------------------------------------
|
||||
|
||||
async def fake_acquire(_session: Any, _context: str) -> None:
|
||||
return None
|
||||
|
||||
async def fake_super_execute(_self: Any, *args: Any, **_kw: Any) -> None:
|
||||
calls.append(args)
|
||||
class _FakeCursor:
|
||||
def __init__(self, recorder: list[Any], raise_exc: Exception | None) -> None:
|
||||
self.recorder: list[Any] = recorder
|
||||
self.raise_exc: Exception | None = raise_exc
|
||||
self.closed: bool = False
|
||||
|
||||
monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire)
|
||||
monkeypatch.setattr(AsyncSession, "execute", fake_super_execute)
|
||||
def execute(self, sql: str, params: Any = None) -> None:
|
||||
if self.raise_exc is not None:
|
||||
raise self.raise_exc
|
||||
self.recorder.append((sql, params))
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _FakeDBAPIConn:
|
||||
def __init__(self, recorder: list[Any], raise_exc: Exception | None = None) -> None:
|
||||
self._cursor: _FakeCursor = _FakeCursor(recorder, raise_exc)
|
||||
|
||||
def cursor(self) -> _FakeCursor:
|
||||
return self._cursor
|
||||
|
||||
|
||||
def test_checkout_hook_sets_application_name_from_request_context() -> None:
|
||||
recorder: list[Any] = []
|
||||
conn = _FakeDBAPIConn(recorder)
|
||||
token = db_module.request_context.set("request:trace-ctx")
|
||||
try:
|
||||
session = db_module.SessionLocal()
|
||||
await session.execute("SELECT 1")
|
||||
db_module._set_application_name_on_checkout(conn, None, None) # pyright: ignore[reportPrivateUsage]
|
||||
finally:
|
||||
db_module.request_context.reset(token)
|
||||
|
||||
# First super().execute is the set_config, then the real statement.
|
||||
assert any("set_config" in str(c[0]) for c in calls)
|
||||
set_config_call = next(c for c in calls if "set_config" in str(c[0]))
|
||||
assert set_config_call[1] == {"name": "request:trace-ctx"}
|
||||
assert len(recorder) == 1
|
||||
sql, params = recorder[0]
|
||||
assert "set_config" in sql and "application_name" in sql
|
||||
assert params == ("request:trace-ctx",)
|
||||
assert conn._cursor.closed is True # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_lifecycle_methods_reset_acquired_flag(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
async def fake_acquire(_session: Any, _context: str) -> None:
|
||||
return None
|
||||
def test_checkout_hook_defaults_to_unknown_without_context() -> None:
|
||||
recorder: list[Any] = []
|
||||
conn = _FakeDBAPIConn(recorder)
|
||||
token = db_module.request_context.set(None)
|
||||
try:
|
||||
db_module._set_application_name_on_checkout(conn, None, None) # pyright: ignore[reportPrivateUsage]
|
||||
finally:
|
||||
db_module.request_context.reset(token)
|
||||
|
||||
async def noop(_self: Any) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire)
|
||||
for method in ("commit", "rollback", "close", "reset"):
|
||||
monkeypatch.setattr(AsyncSession, method, noop)
|
||||
|
||||
session = db_module.SessionLocal()
|
||||
await session.commit() # ensures acquired, commits, then resets
|
||||
assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
# rollback/close/reset must each clear the flag so a reused session
|
||||
# re-acquires (and re-wraps retry) on its next DB use.
|
||||
for method in ("rollback", "close", "reset"):
|
||||
session._honcho_acquired = True # pyright: ignore[reportPrivateUsage]
|
||||
await getattr(session, method)()
|
||||
assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage]
|
||||
assert recorder[0][1] == ("unknown",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_get_and_delete_also_acquire(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The lazy-acquire guarantee covers get/delete, not just execute."""
|
||||
monkeypatch.setattr(settings.DB, "TRACING", False)
|
||||
acquired: list[str] = []
|
||||
|
||||
async def fake_acquire(_session: Any, context: str) -> None:
|
||||
acquired.append(context)
|
||||
|
||||
async def fake_get(_self: Any, *_a: Any, **_k: Any) -> str:
|
||||
return "row"
|
||||
|
||||
async def fake_delete(_self: Any, *_a: Any, **_k: Any) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire)
|
||||
monkeypatch.setattr(AsyncSession, "get", fake_get)
|
||||
monkeypatch.setattr(AsyncSession, "delete", fake_delete)
|
||||
|
||||
session = db_module.SessionLocal()
|
||||
await session.get(object, 1)
|
||||
await session.delete(object())
|
||||
assert acquired == ["unknown"] # acquired once on the first DB-touching call
|
||||
def test_checkout_hook_swallows_errors() -> None:
|
||||
"""A failure tagging the connection must never break the checkout."""
|
||||
conn = _FakeDBAPIConn([], raise_exc=RuntimeError("boom"))
|
||||
# Must not raise.
|
||||
db_module._set_application_name_on_checkout(conn, None, None) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_db_does_not_acquire_at_dependency_entry(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.dependencies import get_db as real_get_db
|
||||
|
||||
acquired: list[str] = []
|
||||
|
||||
async def fake_acquire(_session: Any, context: str) -> None:
|
||||
acquired.append(context)
|
||||
|
||||
monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire)
|
||||
|
||||
dep_gen = real_get_db()
|
||||
await anext(dep_gen)
|
||||
assert acquired == [] # yielding the session must not check out
|
||||
await dep_gen.aclose()
|
||||
assert acquired == [] # finally rollback/close must not check out either
|
||||
# --- deriver polling backoff math --------------------------------------------
|
||||
|
||||
|
||||
def test_polling_backoff_sequence_and_reset(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
@ -250,6 +98,9 @@ def test_polling_backoff_sequence_and_reset(monkeypatch: pytest.MonkeyPatch) ->
|
|||
monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0)
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0)
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 30.0)
|
||||
# Disable jitter so the schedule is asserted exactly (jitter is tested
|
||||
# separately in tests/deriver/test_queue_processing.py::TestPollingJitter).
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0)
|
||||
|
||||
from src.deriver.queue_manager import QueueManager
|
||||
|
||||
|
|
@ -266,6 +117,7 @@ def test_polling_backoff_disabled_stays_constant(
|
|||
) -> None:
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", False)
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0)
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0)
|
||||
|
||||
from src.deriver.queue_manager import QueueManager
|
||||
|
||||
|
|
@ -273,29 +125,6 @@ def test_polling_backoff_disabled_stays_constant(
|
|||
assert [qm._advance_poll_interval() for _ in range(3)] == [1.0, 1.0, 1.0] # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
def test_pool_timeout_must_be_under_retry_budget() -> None:
|
||||
"""A pooled QueuePool checkout that can block past the retry budget is a
|
||||
contradiction (no retry ever happens) and must fail config validation."""
|
||||
from src.config import DBSettings
|
||||
|
||||
# Contradiction: pooled + retry on, but POOL_TIMEOUT >= budget.
|
||||
with pytest.raises(ValueError, match="DB_POOL_TIMEOUT"):
|
||||
DBSettings(
|
||||
POOL_CLASS="default",
|
||||
POOL_TIMEOUT=30,
|
||||
CONNECTION_RETRY_ENABLED=True,
|
||||
CONNECTION_RETRY_MAX_DELAY_SECONDS=10.0,
|
||||
)
|
||||
|
||||
# NullPool has no queue wait, so POOL_TIMEOUT is irrelevant -> allowed.
|
||||
DBSettings(
|
||||
POOL_CLASS="null",
|
||||
POOL_TIMEOUT=30,
|
||||
CONNECTION_RETRY_ENABLED=True,
|
||||
CONNECTION_RETRY_MAX_DELAY_SECONDS=10.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_polling_loop_idle_sleeps_once_per_cycle(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
@ -306,6 +135,7 @@ async def test_polling_loop_idle_sleeps_once_per_cycle(
|
|||
monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0)
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0)
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 8.0)
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0)
|
||||
|
||||
import asyncio
|
||||
|
||||
|
|
|
|||
|
|
@ -40,8 +40,8 @@ class FakeSession:
|
|||
async def test_get_db_yields_lazily_without_checkout_or_tracing(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# get_db must NOT touch the connection or run set_config itself — those now
|
||||
# happen lazily inside HonchoAsyncSession on first DB use, so a handler doing
|
||||
# get_db must NOT touch the connection or run set_config itself — checkout
|
||||
# happens lazily inside the AsyncSession on first DB use, so a handler doing
|
||||
# non-DB work before its first query never pins a connection.
|
||||
fake_db = FakeSession()
|
||||
monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)
|
||||
|
|
|
|||
4
uv.lock
4
uv.lock
|
|
@ -8,7 +8,7 @@ resolution-markers = [
|
|||
]
|
||||
|
||||
[options]
|
||||
exclude-newer = "2026-05-27T18:30:19.790621Z"
|
||||
exclude-newer = "2026-05-28T15:27:36.945866Z"
|
||||
exclude-newer-span = "P5D"
|
||||
|
||||
[manifest]
|
||||
|
|
@ -1159,7 +1159,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "honcho"
|
||||
version = "3.0.8"
|
||||
version = "3.0.9"
|
||||
source = { virtual = "." }
|
||||
dependencies = [
|
||||
{ name = "alembic" },
|
||||
|
|
|
|||
Loading…
Reference in New Issue