wooo/awoooi
1
1
Fork 0
Code Issues 2 Pull Requests 1 Actions 2 Packages Projects Releases Wiki Activity

fix(db): enforce global nullpool session budget

Browse Source
This commit is contained in:
ogt
2026-07-15 15:49:18 +08:00
parent 19d7db93e5
commit f57845ba5e
12 changed files with 709 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

356
apps/api/tests/test_runtime_bootstrap_guards.py
View File

@@ -225,17 +225,11 @@ def test_get_engine_uses_null_pool_for_bounded_background_processes(monkeypatch)
@pytest.mark.asyncio
async def test_null_pool_session_limiter_serializes_worker_db_sessions(
async def test_null_pool_session_limiter_serializes_raw_factory_sessions(
monkeypatch,
) -> None:
from src.db import base as db_base
assert "async with _database_session_slot()" in inspect.getsource(
db_base.get_db
)
assert "async with _database_session_slot()" in inspect.getsource(
db_base.get_db_context
)
monkeypatch.setattr(db_base.settings, "DATABASE_NULL_POOL", True)
monkeypatch.setattr(
db_base.settings,
@@ -245,17 +239,25 @@ async def test_null_pool_session_limiter_serializes_worker_db_sessions(
monkeypatch.setattr(db_base.settings, "DATABASE_POOL_TIMEOUT_SECONDS", 1.0)
monkeypatch.setattr(db_base, "_null_pool_session_limiter", None)
monkeypatch.setattr(db_base, "_null_pool_session_limiter_size", 0)
monkeypatch.setattr(db_base, "_session_factory", None)
monkeypatch.setattr(db_base, "get_engine", lambda: None)
factory = db_base.get_session_factory()
assert factory.class_ is db_base._BudgetedAsyncSession
assert "_database_session_slot" not in inspect.getsource(db_base.get_db)
assert "_database_session_slot" not in inspect.getsource(db_base.get_db_context)
first_entered = asyncio.Event()
release_first = asyncio.Event()
second_entered = asyncio.Event()
async def first_session() -> None:
async with db_base._database_session_slot():
async with factory():
first_entered.set()
await release_first.wait()
async def second_session() -> None:
async with db_base._database_session_slot():
async with factory():
second_entered.set()
first_task = asyncio.create_task(first_session())
@@ -270,6 +272,342 @@ async def test_null_pool_session_limiter_serializes_worker_db_sessions(
assert second_entered.is_set() is True
@pytest.mark.asyncio
async def test_unit_of_work_shares_raw_factory_cap_without_double_acquire(
monkeypatch,
) -> None:
from sqlalchemy.ext.asyncio import AsyncSession
from src.core.unit_of_work import UnitOfWork
from src.db import base as db_base
monkeypatch.setattr(db_base.settings, "DATABASE_NULL_POOL", True)
monkeypatch.setattr(
db_base.settings,
"DATABASE_NULL_POOL_CONCURRENCY_LIMIT",
1,
)
monkeypatch.setattr(db_base.settings, "DATABASE_POOL_TIMEOUT_SECONDS", 1.0)
monkeypatch.setattr(db_base, "_null_pool_session_limiter", None)
monkeypatch.setattr(db_base, "_null_pool_session_limiter_size", 0)
monkeypatch.setattr(db_base, "_session_factory", None)
monkeypatch.setattr(db_base, "get_engine", lambda: None)
executed_session_ids: list[int] = []
async def fake_execute(
session: AsyncSession,
*_args: object,
**_kwargs: object,
) -> object:
executed_session_ids.append(id(session))
return object()
async def fake_commit(_session: AsyncSession) -> None:
return None
async def fake_close(_session: AsyncSession) -> None:
return None
monkeypatch.setattr(AsyncSession, "execute", fake_execute)
monkeypatch.setattr(AsyncSession, "commit", fake_commit)
monkeypatch.setattr(AsyncSession, "close", fake_close)
factory = db_base.get_session_factory()
raw_entered = asyncio.Event()
release_raw = asyncio.Event()
uow_entered = asyncio.Event()
async def raw_factory_session() -> None:
async with factory():
raw_entered.set()
await release_raw.wait()
async def unit_of_work_session() -> None:
async with UnitOfWork(factory, project_id="awoooi"):
uow_entered.set()
raw_task = asyncio.create_task(raw_factory_session())
await raw_entered.wait()
uow_task = asyncio.create_task(unit_of_work_session())
await asyncio.sleep(0)
assert uow_entered.is_set() is False
assert executed_session_ids == []
release_raw.set()
await asyncio.gather(raw_task, uow_task)
assert uow_entered.is_set() is True
assert len(executed_session_ids) == 1
async with asyncio.timeout(0.5):
async with db_base.get_db_context("awoooi"):
pass
assert len(executed_session_ids) == 2
@pytest.mark.asyncio
async def test_unit_of_work_enter_failure_releases_global_session_cap(
monkeypatch,
) -> None:
from sqlalchemy.ext.asyncio import AsyncSession
from src.core.unit_of_work import UnitOfWork
from src.db import base as db_base
monkeypatch.setattr(db_base.settings, "DATABASE_NULL_POOL", True)
monkeypatch.setattr(
db_base.settings,
"DATABASE_NULL_POOL_CONCURRENCY_LIMIT",
1,
)
monkeypatch.setattr(db_base.settings, "DATABASE_POOL_TIMEOUT_SECONDS", 0.2)
monkeypatch.setattr(db_base, "_null_pool_session_limiter", None)
monkeypatch.setattr(db_base, "_null_pool_session_limiter_size", 0)
monkeypatch.setattr(db_base, "_session_factory", None)
monkeypatch.setattr(db_base, "get_engine", lambda: None)
execute_attempts = 0
async def fake_execute(
_session: AsyncSession,
*_args: object,
**_kwargs: object,
) -> object:
nonlocal execute_attempts
execute_attempts += 1
if execute_attempts == 1:
raise RuntimeError("synthetic_uow_enter_failure")
return object()
async def fake_commit(_session: AsyncSession) -> None:
return None
async def fake_close(_session: AsyncSession) -> None:
return None
monkeypatch.setattr(AsyncSession, "execute", fake_execute)
monkeypatch.setattr(AsyncSession, "commit", fake_commit)
monkeypatch.setattr(AsyncSession, "close", fake_close)
factory = db_base.get_session_factory()
with pytest.raises(RuntimeError, match="synthetic_uow_enter_failure"):
async with UnitOfWork(factory, project_id="awoooi"):
pass
async with asyncio.timeout(0.5):
async with UnitOfWork(factory, project_id="awoooi"):
pass
assert execute_attempts == 2
@pytest.mark.asyncio
async def test_budgeted_transaction_enter_failure_releases_global_session_cap(
monkeypatch,
) -> None:
from sqlalchemy.ext.asyncio import AsyncSession, AsyncSessionTransaction
from src.db import base as db_base
monkeypatch.setattr(db_base.settings, "DATABASE_NULL_POOL", True)
monkeypatch.setattr(
db_base.settings,
"DATABASE_NULL_POOL_CONCURRENCY_LIMIT",
1,
)
monkeypatch.setattr(db_base.settings, "DATABASE_POOL_TIMEOUT_SECONDS", 0.2)
monkeypatch.setattr(db_base, "_null_pool_session_limiter", None)
monkeypatch.setattr(db_base, "_null_pool_session_limiter_size", 0)
monkeypatch.setattr(db_base, "_session_factory", None)
monkeypatch.setattr(db_base, "get_engine", lambda: None)
async def fail_start(
_transaction: AsyncSessionTransaction,
is_ctxmanager: bool = False,
) -> AsyncSessionTransaction:
del is_ctxmanager
raise RuntimeError("synthetic_transaction_enter_failure")
async def fake_close(_session: AsyncSession) -> None:
return None
monkeypatch.setattr(AsyncSessionTransaction, "start", fail_start)
monkeypatch.setattr(AsyncSession, "close", fake_close)
factory = db_base.get_session_factory()
with pytest.raises(RuntimeError, match="synthetic_transaction_enter_failure"):
async with factory.begin():
pass
async with asyncio.timeout(0.5):
async with factory():
pass
@pytest.mark.asyncio
async def test_cancelled_transaction_entry_releases_global_session_cap(
monkeypatch,
) -> None:
from sqlalchemy.ext.asyncio import AsyncSession, AsyncSessionTransaction
from src.db import base as db_base
monkeypatch.setattr(db_base.settings, "DATABASE_NULL_POOL", True)
monkeypatch.setattr(
db_base.settings,
"DATABASE_NULL_POOL_CONCURRENCY_LIMIT",
1,
)
monkeypatch.setattr(db_base.settings, "DATABASE_POOL_TIMEOUT_SECONDS", 0.2)
monkeypatch.setattr(db_base, "_null_pool_session_limiter", None)
monkeypatch.setattr(db_base, "_null_pool_session_limiter_size", 0)
monkeypatch.setattr(db_base, "_session_factory", None)
monkeypatch.setattr(db_base, "get_engine", lambda: None)
transaction_start_entered = asyncio.Event()
never_complete = asyncio.Event()
async def block_start(
_transaction: AsyncSessionTransaction,
is_ctxmanager: bool = False,
) -> AsyncSessionTransaction:
del is_ctxmanager
transaction_start_entered.set()
await never_complete.wait()
raise AssertionError("unreachable")
async def fake_close(_session: AsyncSession) -> None:
return None
monkeypatch.setattr(AsyncSessionTransaction, "start", block_start)
monkeypatch.setattr(AsyncSession, "close", fake_close)
factory = db_base.get_session_factory()
async def enter_transaction() -> None:
async with factory.begin():
pass
transaction_task = asyncio.create_task(enter_transaction())
await transaction_start_entered.wait()
transaction_task.cancel()
with pytest.raises(asyncio.CancelledError):
await transaction_task
async with asyncio.timeout(0.5):
async with factory():
pass
@pytest.mark.asyncio
async def test_cancelled_slot_waiter_does_not_consume_or_over_release_capacity(
monkeypatch,
) -> None:
from sqlalchemy.ext.asyncio import AsyncSession
from src.db import base as db_base
monkeypatch.setattr(db_base.settings, "DATABASE_NULL_POOL", True)
monkeypatch.setattr(
db_base.settings,
"DATABASE_NULL_POOL_CONCURRENCY_LIMIT",
1,
)
monkeypatch.setattr(db_base.settings, "DATABASE_POOL_TIMEOUT_SECONDS", 1.0)
monkeypatch.setattr(db_base, "_null_pool_session_limiter", None)
monkeypatch.setattr(db_base, "_null_pool_session_limiter_size", 0)
monkeypatch.setattr(db_base, "_session_factory", None)
monkeypatch.setattr(db_base, "get_engine", lambda: None)
async def fake_close(_session: AsyncSession) -> None:
return None
monkeypatch.setattr(AsyncSession, "close", fake_close)
factory = db_base.get_session_factory()
first_entered = asyncio.Event()
release_first = asyncio.Event()
async def hold_first_slot() -> None:
async with factory():
first_entered.set()
await release_first.wait()
async def wait_for_slot() -> None:
async with factory():
pass
first_task = asyncio.create_task(hold_first_slot())
await first_entered.wait()
cancelled_waiter = asyncio.create_task(wait_for_slot())
await asyncio.sleep(0)
cancelled_waiter.cancel()
with pytest.raises(asyncio.CancelledError):
await cancelled_waiter
release_first.set()
await first_task
async with asyncio.timeout(0.5):
async with factory():
pass
@pytest.mark.asyncio
async def test_unit_of_work_reentry_fails_without_replacing_active_session(
monkeypatch,
) -> None:
from sqlalchemy.ext.asyncio import AsyncSession
from src.core.unit_of_work import UnitOfWork
from src.db import base as db_base
monkeypatch.setattr(db_base.settings, "DATABASE_NULL_POOL", True)
monkeypatch.setattr(
db_base.settings,
"DATABASE_NULL_POOL_CONCURRENCY_LIMIT",
1,
)
monkeypatch.setattr(db_base.settings, "DATABASE_POOL_TIMEOUT_SECONDS", 0.2)
monkeypatch.setattr(db_base, "_null_pool_session_limiter", None)
monkeypatch.setattr(db_base, "_null_pool_session_limiter_size", 0)
monkeypatch.setattr(db_base, "_session_factory", None)
monkeypatch.setattr(db_base, "get_engine", lambda: None)
async def fake_execute(
_session: AsyncSession,
*_args: object,
**_kwargs: object,
) -> object:
return object()
async def fake_commit(_session: AsyncSession) -> None:
return None
async def fake_close(_session: AsyncSession) -> None:
return None
monkeypatch.setattr(AsyncSession, "execute", fake_execute)
monkeypatch.setattr(AsyncSession, "commit", fake_commit)
monkeypatch.setattr(AsyncSession, "close", fake_close)
factory = db_base.get_session_factory()
unit_of_work = UnitOfWork(factory, project_id="awoooi")
async with unit_of_work:
active_session = unit_of_work.session
with pytest.raises(
RuntimeError,
match="UnitOfWork cannot be entered more than once",
):
await unit_of_work.__aenter__()
assert unit_of_work.session is active_session
async with asyncio.timeout(0.5):
async with factory():
pass
@pytest.mark.asyncio
async def test_init_db_serializes_bootstrap_ddl_with_advisory_lock(monkeypatch):
from src.db import base as db_base