From 7b0ade398741aada5f58f9712892a9fe5a21e36b Mon Sep 17 00:00:00 2001 From: Jeroen Ubbink Date: Fri, 24 Jul 2026 22:26:25 +0200 Subject: [PATCH] fix: stop AsyncSession pool exhaustion in balance reads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fetch_all_balances shared a single AsyncSession across the tasks it ran with asyncio.gather. AsyncSession is not safe for concurrent use: the overlapping queries raised "concurrent operations are not permitted" and left connections wedged until the QueuePool was exhausted, after which admin/API endpoints returned 500/502 until a container restart. Read all outstanding user liabilities up front in one short-lived session with a single grouped query (balances_by_mint_and_unit), then run the per-mint balance checks concurrently with no session in scope. A failure reading liabilities now degrades gracefully — the page still reports the known wallet custody and blanks only the unknowable user/owner split, tagging each mint with the error — instead of 500-ing the whole page. periodic_payout reads each liability fresh, immediately before the payout decision, rather than from a single pre-loop snapshot: the per-mint round trip is slow, and a user top-up during the cycle would otherwise let a later mint/unit act on a stale-low liability and over-send funds owed to users (related to the payout-safety concern in issue #611). Co-Authored-By: Claude Opus 4.8 --- routstr/core/db.py | 35 ++- routstr/wallet.py | 92 +++++--- tests/unit/test_balances_by_mint_and_unit.py | 100 +++++++++ tests/unit/test_fetch_all_balances.py | 219 ++++++++++++++++++- tests/unit/test_periodic_payout.py | 61 +++++- 5 files changed, 464 insertions(+), 43 deletions(-) create mode 100644 tests/unit/test_balances_by_mint_and_unit.py diff --git a/routstr/core/db.py b/routstr/core/db.py index 16cfedfd..ab263be4 100644 --- a/routstr/core/db.py +++ b/routstr/core/db.py @@ -717,14 +717,37 @@ async def complete_routstr_fee_payout( return result.rowcount == 1 -async def balances_for_mint_and_unit( - db_session: AsyncSession, mint_url: str, unit: str -) -> int: - query = select(func.sum(ApiKey.balance)).where( - ApiKey.refund_mint_url == mint_url, ApiKey.refund_currency == unit +async def balances_by_mint_and_unit( + db_session: AsyncSession, mint_urls: list[str], units: list[str] +) -> dict[tuple[str, str], int]: + """Sum outstanding user balances (msats) grouped by (mint_url, unit). + + One query for every requested mint/unit pair. Pairs with no matching rows + are absent from the mapping — callers should default those to 0. + """ + if not mint_urls or not units: + return {} + query = ( + select( + ApiKey.refund_mint_url, + ApiKey.refund_currency, + func.sum(ApiKey.balance), + ) + .where( + col(ApiKey.refund_mint_url).in_(mint_urls), + col(ApiKey.refund_currency).in_(units), + ) + .group_by(col(ApiKey.refund_mint_url), col(ApiKey.refund_currency)) ) result = await db_session.exec(query) - return result.one() or 0 + # IN(...) filters out NULL mint_url/currency at the SQL level, so the group + # keys are never None at runtime; the guard also narrows the nullable + # columns from ``str | None`` to ``str`` for the typed return. + return { + (mint_url, unit): total or 0 + for mint_url, unit, total in result.all() + if mint_url is not None and unit is not None + } async def init_db() -> None: diff --git a/routstr/wallet.py b/routstr/wallet.py index dd92d913..99b77e5c 100644 --- a/routstr/wallet.py +++ b/routstr/wallet.py @@ -837,7 +837,7 @@ async def fetch_all_balances( units = ["sat", "msat"] async def fetch_balance( - session: db.AsyncSession, mint_url: str, unit: str + mint_url: str, unit: str, user_balance: int ) -> BalanceDetail: try: wallet = await get_wallet(mint_url, unit) @@ -845,7 +845,6 @@ async def fetch_all_balances( wallet, mint_url, unit, not_reserved=True ) proofs = await slow_filter_spend_proofs(proofs, wallet) - user_balance = await db.balances_for_mint_and_unit(session, mint_url, unit) if unit == "sat": user_balance = user_balance // 1000 proofs_balance = sum(proof.amount for proof in proofs) @@ -878,40 +877,66 @@ async def fetch_all_balances( if settings.primary_mint and settings.primary_mint not in mint_urls: mint_urls.append(settings.primary_mint) - # Create tasks for all mint/unit combinations - async with db.create_session() as session: - tasks = [ - fetch_balance(session, mint_url, unit) - for mint_url in mint_urls - for unit in units - ] + # Read all outstanding user liabilities up front in one short-lived DB + # session and a single grouped query. AsyncSession is not safe for + # concurrent use, so the session must NOT be shared across the gathered + # mint checks below (that raises "concurrent operations are not permitted" + # and can wedge connections until the pool is exhausted). A DB failure here + # must still degrade gracefully rather than 500 the whole balances page. + user_balances: dict[tuple[str, str], int] = {} + liabilities_error: str | None = None + try: + async with db.create_session() as session: + user_balances = await db.balances_by_mint_and_unit( + session, mint_urls, units + ) + except Exception as e: + logger.error(f"Error reading user balances: {e}") + liabilities_error = str(e) - # Run all tasks concurrently - balance_details = list(await asyncio.gather(*tasks)) + # Run the per-mint balance checks concurrently — no DB session involved. + tasks = [ + fetch_balance(mint_url, unit, user_balances.get((mint_url, unit), 0)) + for mint_url in mint_urls + for unit in units + ] + balance_details = list(await asyncio.gather(*tasks)) - # Calculate totals + # Compute totals BEFORE tagging any liability-read failure. A failed + # liability read does not invalidate custody, so the known wallet balance + # must still be summed; only the per-user split is unknowable. A per-mint + # fetch failure (``error`` already set inside fetch_balance) means custody + # for that mint is genuinely unknown, so those details are skipped. total_wallet_balance_sats = 0 total_user_balance_sats = 0 - for detail in balance_details: - if not detail.get("error"): - # Convert to sats for total calculation - unit = detail["unit"] - proofs_balance_sats = ( - detail["wallet_balance"] - if unit == "sat" - else detail["wallet_balance"] // 1000 - ) - user_balance_sats = ( + if detail.get("error"): + continue + unit = detail["unit"] + total_wallet_balance_sats += ( + detail["wallet_balance"] + if unit == "sat" + else detail["wallet_balance"] // 1000 + ) + if liabilities_error is None: + total_user_balance_sats += ( detail["user_balance"] if unit == "sat" else detail["user_balance"] // 1000 ) - total_wallet_balance_sats += proofs_balance_sats - total_user_balance_sats += user_balance_sats - - owner_balance = total_wallet_balance_sats - total_user_balance_sats + if liabilities_error is None: + owner_balance = total_wallet_balance_sats - total_user_balance_sats + else: + # Liabilities are unknown: report custody truthfully but never claim any + # of it as owner profit (owner = wallet - unknown liabilities). Surface + # the failure on each detail without discarding a more specific per-mint + # error, and blank the unknowable per-user/owner split. + owner_balance = 0 + for detail in balance_details: + detail["user_balance"] = 0 + detail["owner_balance"] = 0 + detail.setdefault("error", liabilities_error) return ( balance_details, @@ -935,9 +960,10 @@ async def periodic_payout() -> None: if settings.primary_mint and settings.primary_mint not in mint_urls: mint_urls.append(settings.primary_mint) + units = ["sat", "msat"] async with db.create_session() as session: for mint_url in mint_urls: - for unit in ["sat", "msat"]: + for unit in units: # Isolate failures per mint/unit so one slow or failing # mint does not abort payout for every other mint/unit. try: @@ -947,9 +973,17 @@ async def periodic_payout() -> None: ) proofs = await slow_filter_spend_proofs(proofs, wallet) await asyncio.sleep(5) - user_balance = await db.balances_for_mint_and_unit( - session, mint_url, unit + # Read the liability fresh, right before deciding the + # payout. The mint round-trip above is slow; reading + # it once before the loop would let a user top-up + # during the cycle go unseen, so a later (mint, unit) + # would act on a stale-low liability and over-send + # funds owed to users. The loop is sequential, so + # reusing this session per iteration is safe. + balances = await db.balances_by_mint_and_unit( + session, [mint_url], [unit] ) + user_balance = balances.get((mint_url, unit), 0) if unit == "sat": user_balance = user_balance // 1000 proofs_balance = sum(proof.amount for proof in proofs) diff --git a/tests/unit/test_balances_by_mint_and_unit.py b/tests/unit/test_balances_by_mint_and_unit.py new file mode 100644 index 00000000..bb1b3c81 --- /dev/null +++ b/tests/unit/test_balances_by_mint_and_unit.py @@ -0,0 +1,100 @@ +"""Real-DB coverage for db.balances_by_mint_and_unit. + +Verifies the grouped liability query used by fetch_all_balances: it sums +balances per (mint_url, unit), filters to the requested mints/units, excludes +NULL mint/currency rows, and returns nothing for empty inputs. +""" + +from typing import AsyncGenerator + +import pytest +from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine +from sqlalchemy.pool import StaticPool +from sqlmodel import SQLModel +from sqlmodel.ext.asyncio.session import AsyncSession + +from routstr.core.db import ApiKey, balances_by_mint_and_unit + + +def _make_engine() -> AsyncEngine: + return create_async_engine( + "sqlite+aiosqlite://", + poolclass=StaticPool, + connect_args={"check_same_thread": False}, + ) + + +@pytest.fixture +async def session() -> "AsyncGenerator[AsyncSession, None]": + engine = _make_engine() + async with engine.begin() as conn: + await conn.run_sync(SQLModel.metadata.create_all) + db_session = AsyncSession(engine, expire_on_commit=False) + try: + yield db_session + finally: + await db_session.close() + await engine.dispose() + + +async def _add_key( + session: AsyncSession, + hashed_key: str, + balance: int, + mint_url: str | None, + currency: str | None, +) -> None: + session.add( + ApiKey( + hashed_key=hashed_key, + balance=balance, + refund_mint_url=mint_url, + refund_currency=currency, + ) + ) + await session.commit() + + +@pytest.mark.asyncio +async def test_sums_and_groups_by_mint_and_unit(session: AsyncSession) -> None: + await _add_key(session, "a", 1000, "http://m1", "sat") + await _add_key(session, "b", 500, "http://m1", "sat") + await _add_key(session, "c", 7000, "http://m1", "msat") + await _add_key(session, "d", 200, "http://m2", "sat") + + result = await balances_by_mint_and_unit( + session, ["http://m1", "http://m2"], ["sat", "msat"] + ) + + assert result[("http://m1", "sat")] == 1500 + assert result[("http://m1", "msat")] == 7000 + assert result[("http://m2", "sat")] == 200 + + +@pytest.mark.asyncio +async def test_filters_out_unrequested_mints_and_units(session: AsyncSession) -> None: + await _add_key(session, "a", 1000, "http://wanted", "sat") + await _add_key(session, "b", 999, "http://other", "sat") + await _add_key(session, "c", 888, "http://wanted", "usd") + + result = await balances_by_mint_and_unit(session, ["http://wanted"], ["sat"]) + + assert result == {("http://wanted", "sat"): 1000} + + +@pytest.mark.asyncio +async def test_excludes_rows_with_null_mint_or_currency(session: AsyncSession) -> None: + await _add_key(session, "a", 1000, "http://m1", "sat") + await _add_key(session, "b", 4242, None, None) + + result = await balances_by_mint_and_unit(session, ["http://m1"], ["sat"]) + + assert result == {("http://m1", "sat"): 1000} + + +@pytest.mark.asyncio +async def test_empty_inputs_return_empty_mapping(session: AsyncSession) -> None: + await _add_key(session, "a", 1000, "http://m1", "sat") + + assert await balances_by_mint_and_unit(session, [], ["sat"]) == {} + assert await balances_by_mint_and_unit(session, ["http://m1"], []) == {} diff --git a/tests/unit/test_fetch_all_balances.py b/tests/unit/test_fetch_all_balances.py index 433e84d4..977c2100 100644 --- a/tests/unit/test_fetch_all_balances.py +++ b/tests/unit/test_fetch_all_balances.py @@ -1,4 +1,7 @@ -from contextlib import asynccontextmanager +import asyncio +from collections.abc import AsyncIterator, Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -11,9 +14,70 @@ async def _fake_session(): # type: ignore[no-untyped-def] yield MagicMock() -def _patches( # type: ignore[no-untyped-def] - proof_amount: int = 1000, user_balance_msats: int = 0 -): +class _TrackingSession: + """Stand-in for ``AsyncSession`` that records concurrent use. + + Real ``AsyncSession`` is not safe for concurrent use: if a second + coroutine issues a query while an ``await``-ed query is still in flight + on the same session it raises ``"This session is provisioning a new + connection; concurrent operations are not permitted"`` and can leave a + connection wedged. + + Every awaited query method (``exec``, ``execute``, ``scalar``, ``get`` …) + is resolved through ``__getattr__`` to the same tracked coroutine, so the + double detects concurrent use regardless of which session API the code + under test happens to call — it counts how many coroutines are inside a + query at once and remembers the peak. + """ + + def __init__(self) -> None: + self._in_flight = 0 + self.max_concurrency = 0 + self.query_calls = 0 + + def __getattr__(self, name: str): # type: ignore[no-untyped-def] + async def _tracked(*args: object, **kwargs: object) -> MagicMock: + self.query_calls += 1 + self._in_flight += 1 + self.max_concurrency = max(self.max_concurrency, self._in_flight) + try: + # Yield to the event loop so any concurrently-scheduled task + # sharing this session can enter a query before we exit. + await asyncio.sleep(0.01) + result = MagicMock() + result.one.return_value = 0 + result.all.return_value = [] + return result + finally: + self._in_flight -= 1 + + return _tracked + + +def _tracking_session_factory() -> tuple[ + Callable[[], AbstractAsyncContextManager[_TrackingSession]], + list[_TrackingSession], +]: + """Patch replacement for ``db.create_session`` that records every session. + + Each ``async with db.create_session()`` gets a fresh ``_TrackingSession``, + appended to ``sessions`` so the test can assert that *no single* session + was ever used concurrently — true whether the fix precomputes balances in + one query (one session, used serially) or opens a session per task. + """ + sessions: list[_TrackingSession] = [] + + @asynccontextmanager + async def _factory() -> "AsyncIterator[_TrackingSession]": + session = _TrackingSession() + sessions.append(session) + yield session + + return _factory, sessions + + +def _mint_patches(proof_amount: int = 1000): # type: ignore[no-untyped-def] + """Patches for the mint/proof side of ``fetch_all_balances`` (no DB).""" proof = MagicMock(amount=proof_amount) return [ patch("routstr.wallet.get_wallet", AsyncMock(return_value=MagicMock())), @@ -25,9 +89,21 @@ def _patches( # type: ignore[no-untyped-def] "routstr.wallet.slow_filter_spend_proofs", AsyncMock(side_effect=lambda proofs, wallet: proofs), ), + ] + + +def _patches( # type: ignore[no-untyped-def] + proof_amount: int = 1000, user_balance_msats: int = 0 +): + async def _grouped( # type: ignore[no-untyped-def] + session: object, mint_urls: list[str], units: list[str] + ): + return {(m, u): user_balance_msats for m in mint_urls for u in units} + + return _mint_patches(proof_amount) + [ patch( - "routstr.wallet.db.balances_for_mint_and_unit", - AsyncMock(return_value=user_balance_msats), + "routstr.wallet.db.balances_by_mint_and_unit", + AsyncMock(side_effect=_grouped), ), patch("routstr.wallet.db.create_session", _fake_session), ] @@ -79,6 +155,137 @@ async def test_fetch_all_balances_reports_liability_when_wallet_is_empty() -> No assert owner == -5 +@pytest.mark.asyncio +async def test_fetch_all_balances_does_not_use_one_session_concurrently() -> None: + """A single AsyncSession must never be shared across the gathered tasks. + + ``fetch_all_balances`` fans out one balance lookup per (mint, unit) with + ``asyncio.gather``. Passing one session into all of them makes concurrent + coroutines issue queries on the same AsyncSession — unsafe, and the source + of the observed "concurrent operations are not permitted" error that + cascades into connection-pool exhaustion. + + We drive it with several mint/unit combinations and a session whose + ``exec`` records peak concurrency. If any session is entered by more than + one coroutine at a time the finding is present. + """ + from routstr.core.settings import settings + + factory, sessions = _tracking_session_factory() + + # NB: db.balances_by_mint_and_unit is intentionally NOT patched here — we + # let the real helper run so it issues a query on the session, which is + # what exercises (and detects) concurrent session use. + patches = _mint_patches() + [ + patch("routstr.wallet.db.create_session", factory), + ] + + with patch.object( + settings, "cashu_mints", ["http://mint-a:3338", "http://mint-b:3338"] + ), patch.object(settings, "primary_mint", "http://mint-a:3338"): + for p in patches: + p.start() + try: + await fetch_all_balances(units=["sat", "msat"]) + finally: + patch.stopall() + + # The liabilities must be read in exactly ONE short-lived session with ONE + # query, before the concurrent mint fan-out — not a session per gathered + # task, and never a session shared across tasks. Asserting the exact shape + # (rather than a peak-concurrency counter, which cannot trip once no session + # crosses the gather) makes the guarantee explicit and catches a regression + # that reintroduces per-task DB access. + assert len(sessions) == 1, ( + f"expected exactly one create_session() for the up-front liability read, " + f"got {len(sessions)}" + ) + assert sessions[0].query_calls == 1, ( + f"expected a single grouped liability query, got " + f"{sessions[0].query_calls} — per-(mint,unit) querying has returned" + ) + # And that single session was never entered concurrently. + assert sessions[0].max_concurrency <= 1 + + +@pytest.mark.asyncio +async def test_fetch_all_balances_degrades_when_liability_read_fails() -> None: + """A DB failure reading liabilities must not 500 the whole balances page. + + A failed liability read does not invalidate custody, so the *known* wallet + balance is still reported (both per-mint and in the wallet total). Only the + unknowable per-user/owner split is blanked, and owner is reported as 0 so + the custody is never claimed as owner profit. + """ + from routstr.core.settings import settings + + patches = _mint_patches(proof_amount=1000) + [ + patch( + "routstr.wallet.db.balances_by_mint_and_unit", + AsyncMock(side_effect=RuntimeError("db pool exhausted")), + ), + patch("routstr.wallet.db.create_session", _fake_session), + ] + + with patch.object(settings, "cashu_mints", []), patch.object( + settings, "primary_mint", "http://primary:3338" + ): + for p in patches: + p.start() + try: + details, total_wallet, total_user, owner = await fetch_all_balances( + units=["sat"] + ) + finally: + patch.stopall() + + assert details[0]["error"] == "db pool exhausted" + assert details[0]["user_balance"] == 0 + assert details[0]["owner_balance"] == 0 + # Custody is known and must not be hidden by the liability-read failure. + assert details[0]["wallet_balance"] == 1000 + assert total_wallet == 1000 + # The user/owner split is unknown: blanked, and never claimed as profit. + assert total_user == 0 + assert owner == 0 + + +@pytest.mark.asyncio +async def test_liability_error_does_not_clobber_a_specific_mint_error() -> None: + """A per-mint fetch failure keeps its specific error under a liability error. + + When the up-front liability read fails, every detail is tagged so the UI + shows liabilities are unknown — but a detail that already failed at the mint + level carries a more specific message, which must not be overwritten by the + generic liability-read error. + """ + from routstr.core.settings import settings + + patches: list[Any] = [ + patch( + "routstr.wallet.get_wallet", + AsyncMock(side_effect=RuntimeError("mint down")), + ), + patch( + "routstr.wallet.db.balances_by_mint_and_unit", + AsyncMock(side_effect=RuntimeError("db pool exhausted")), + ), + patch("routstr.wallet.db.create_session", _fake_session), + ] + + with patch.object(settings, "cashu_mints", []), patch.object( + settings, "primary_mint", "http://primary:3338" + ): + for p in patches: + p.start() + try: + details, *_ = await fetch_all_balances(units=["sat"]) + finally: + patch.stopall() + + assert details[0]["error"] == "mint down" + + @pytest.mark.asyncio async def test_fetch_all_balances_no_duplicate_primary_mint() -> None: """primary_mint already in cashu_mints is not inspected twice.""" diff --git a/tests/unit/test_periodic_payout.py b/tests/unit/test_periodic_payout.py index 54ebadb4..0ce944b6 100644 --- a/tests/unit/test_periodic_payout.py +++ b/tests/unit/test_periodic_payout.py @@ -74,7 +74,7 @@ async def test_periodic_payout_includes_primary_mint_not_in_cashu_mints() -> Non "routstr.wallet.slow_filter_spend_proofs", AsyncMock(side_effect=lambda proofs, wallet: proofs), ), patch( - "routstr.wallet.db.balances_for_mint_and_unit", AsyncMock(return_value=0) + "routstr.wallet.db.balances_by_mint_and_unit", AsyncMock(return_value={}) ), patch("routstr.wallet.raw_send_to_lnurl", raw_send): with pytest.raises(_LoopBreak): await periodic_payout() @@ -112,7 +112,7 @@ async def test_periodic_payout_isolates_failing_mint() -> None: "routstr.wallet.slow_filter_spend_proofs", AsyncMock(side_effect=lambda proofs, wallet: proofs), ), patch( - "routstr.wallet.db.balances_for_mint_and_unit", AsyncMock(return_value=0) + "routstr.wallet.db.balances_by_mint_and_unit", AsyncMock(return_value={}) ), patch("routstr.wallet.raw_send_to_lnurl", raw_send): with pytest.raises(_LoopBreak): await periodic_payout() @@ -126,6 +126,63 @@ async def test_periodic_payout_isolates_failing_mint() -> None: assert raw_send.await_count == 2 # good mint paid for both units +@pytest.mark.asyncio +async def test_periodic_payout_reads_liability_fresh_per_iteration() -> None: + """The liability must be read fresh inside the loop, after the mint round-trip. + + Reading all liabilities once *before* the loop lets a user top-up during the + slow, per-mint payout cycle go unseen: a later (mint, unit) then computes + ``available_balance = fresh_proofs - stale_liability`` and can over-send + funds owed to users. The read must happen per (mint, unit), after the inner + ``asyncio.sleep(5)``, so it reflects the balance at payout time. + """ + from routstr.core.settings import settings + + events: list[str] = [] + + async def _sleep(seconds: float) -> None: + if seconds == _INTERVAL: + events.append("interval") + if events.count("interval") >= 2: + raise _LoopBreak() + else: + events.append(f"sleep{int(seconds)}") + + async def _liability( + session: Any, mint_urls: list[str], units: list[str] + ) -> dict[tuple[str, str], int]: + events.append("liability") + return {} + + with patch.object(settings, "cashu_mints", ["http://m:3338"]), patch.object( + settings, "primary_mint", "http://m:3338" + ), patch.object(settings, "receive_ln_address", "owner@ln.tld"), patch.object( + settings, "payout_interval_seconds", _INTERVAL + ), patch.object(settings, "min_payout_sat", 10), patch( + "routstr.wallet.asyncio.sleep", _sleep + ), patch("routstr.wallet.db.create_session", _fake_session), patch( + "routstr.wallet.get_wallet", AsyncMock(return_value=MagicMock()) + ), patch( + "routstr.wallet.get_proofs_per_mint_and_unit", + MagicMock(return_value=[MagicMock(amount=100_000)]), + ), patch( + "routstr.wallet.slow_filter_spend_proofs", + AsyncMock(side_effect=lambda proofs, wallet: proofs), + ), patch( + "routstr.wallet.db.balances_by_mint_and_unit", + AsyncMock(side_effect=_liability), + ), patch("routstr.wallet.raw_send_to_lnurl", AsyncMock(return_value=1000)): + with pytest.raises(_LoopBreak): + await periodic_payout() + + # One fresh liability read per (mint, unit) pair (sat + msat), not a single + # pre-loop snapshot shared across the whole cycle. + assert events.count("liability") == 2 + # And the first read happens only after the first inner mint sleep, proving + # it is read inside the loop at payout time rather than before it. + assert events.index("liability") > events.index("sleep5") + + @pytest.mark.asyncio async def test_periodic_payout_handles_session_creation_failure() -> None: """A db.create_session failure is logged and the payout loop continues."""