Procházet zdrojové kódy

fix: close SQLite connections on exit + defer WS balance snapshot to background pool

Fix 1 (critical): get_connection() was returning a raw sqlite3.Connection.
The 'with conn:' pattern only commits — it does not close. Connections
accumulated for the entire process lifetime, each holding pager cache
memory, causing 3GB RAM usage over time. Now a @contextmanager that
guarantees conn.close() in finally.

Fix 2: _record_order_status_update called capture_account_balance_snapshot
synchronously when driven by a WS order-finish message. This blocked the
event loop for HTTP calls + per-asset price lookups. Now WS-driven calls
(account_id is None) defer the snapshot to a single-worker
ThreadPoolExecutor, serialising snapshots without blocking the WS loop.
Tool-driven calls (query_order, cancel_order, etc.) are unchanged.

Added test for WS-path deferral. All 19 tests pass.
Remaining fixes documented in MEMORY_FIXES_PLAN.md.
Lukas Goldschmidt před 1 týdnem
rodič
revize
547c82026b

+ 86 - 0
MEMORY_FIXES_PLAN.md

@@ -0,0 +1,86 @@
+# Memory Leak Remediation Plan — Remaining Fixes
+
+Applied fixes (this commit):
+- **Fix 1**: SQLite connection leak — `get_connection()` now a contextmanager that closes on exit (`storage.py`)
+- **Fix 2**: WS handler no longer blocks on balance snapshot — deferred to single-worker ThreadPoolExecutor (`services_orders.py`)
+
+---
+
+## Fix 3 — WS reconnect backoff (`bitstamp_private_ws.py`)
+
+**Problem**: `private_ws_main` retries with a fixed 5-second delay after any exception. If Bitstamp drops the connection repeatedly (rate limit, network blip), this creates a tight reconnect loop. Each reconnect calls `list_accounts` (opens a DB connection) and `_get_token` (HTTP call to Bitstamp), amplifying load.
+
+**Fix**: Replace the fixed `WS_RECONNECT_SECONDS` sleep with exponential backoff:
+- Start at 5s, double on each consecutive failure, cap at 60s.
+- Reset to 5s after a successful connection that lasts > 30 seconds.
+- Log the backoff delay so tight loops are visible in logs.
+
+**Location**: `private_ws_main` and `_run_once` in `bitstamp_private_ws.py`.
+
+---
+
+## Fix 4 — `_CACHE_LOCKS` unbounded growth (`services_bitstamp.py`)
+
+**Problem**: `_CACHE_LOCKS` dict grows every time a new `cache_key` is used (per account, per market symbol, etc.). Locks are never removed, even when the account is deleted or the cache entry expires.
+
+**Fix options** (pick one):
+- **A**: Use `weakref.WeakValueDictionary` — locks get GC'd when no caller holds a reference. Simplest, but the `with _cache_lock(cache_key):` pattern keeps a strong reference during use, so the lock survives the call and is collected after.
+- **B**: Add a `_remove_cache_lock(cache_key)` call in `invalidate_account_cache` and account deletion paths. More explicit but requires wiring up cleanup at every deletion site.
+
+**Location**: `_CACHE_LOCKS`, `_cache_lock` in `services_bitstamp.py:24-46`.
+
+---
+
+## Fix 5 — `api_cache` table eviction (`repo.py` / `server.py`)
+
+**Problem**: `cache_put` does an upsert, so existing keys don't duplicate. But expired rows are never deleted — `cache_get` skips them, but they stay in the SQLite file. Over weeks, the table grows and the file bloats. Same applies to `order_records` and `account_balance_snapshots` (no retention policy).
+
+**Fix**: Add a `_evict_expired_cache()` function in `repo.py`:
+```python
+def evict_expired_cache() -> int:
+    now = utc_now_iso()
+    with get_connection() as conn:
+        count = conn.execute("DELETE FROM api_cache WHERE expires_at <= ?", (now,)).rowcount
+        conn.commit()
+    return count
+```
+Call it from the existing metadata refresh loop (`_metadata_refresh_loop` in `server.py`) every cycle — it already runs every 24 hours, which is sufficient for cache eviction.
+
+For `order_records` and `account_balance_snapshots`, consider a configurable retention (e.g., keep last 90 days, or last 10,000 rows per account). This is a bigger change and should be a separate discussion.
+
+**Location**: `repo.py`, `_metadata_refresh_loop` in `server.py`.
+
+---
+
+## Fix 6 — Price lookup caching in `crypto_client.py`
+
+**Problem**: `_price_for_asset` calls `get_crypto_price` which calls `get_price_sync`, which creates a `ThreadPoolExecutor` + `asyncio.run()` + MCP SSE connection **per asset, per call**. A balance with 10 assets triggers 10 separate SSE connections. Now that this runs in the background (Fix 2), it won't block the WS loop, but it's still very wasteful.
+
+**Fix**: Add a short-TTL in-memory price cache (e.g., 5 seconds) in `crypto_client.py`:
+```python
+_PRICE_CACHE: dict[str, tuple[float, float]] = {}  # symbol -> (price, expires_at)
+_PRICE_CACHE_TTL = 5.0
+
+def get_price_sync(symbol: str, base_url: str | None = None) -> dict[str, Any]:
+    now = time.monotonic()
+    cached = _PRICE_CACHE.get(symbol)
+    if cached and cached[1] > now:
+        return {"symbol": symbol.upper(), "price": cached[0]}
+    # ... existing fetch logic ...
+    if "price" in result:
+        _PRICE_CACHE[symbol] = (float(result["price"]), now + _PRICE_CACHE_TTL)
+    return result
+```
+
+This caps price lookups to at most 1 per 5 seconds per symbol, regardless of how many balance snapshots fire. Freshness is preserved (5s TTL is short enough for downstream services that poll balance info).
+
+**Location**: `crypto_client.py`, `_price_for_asset` in `services_bitstamp.py`.
+
+---
+
+## Deployment & Monitoring
+
+1. After deploying Fixes 1+2, monitor container RSS — should stabilize at 50-150 MB instead of climbing to GBs.
+2. Check logs for "private websocket supervisor failed" frequency — if reconnects are frequent, Fix 3 is needed urgently.
+3. Check SQLite file size (`data/exec_mcp.sqlite3`) — if growing, Fix 5 is needed.
+4. Check `docker stats` for the container after 24-48 hours of running to confirm memory is flat.

+ 35 - 4
src/exec_mcp/services_orders.py

@@ -5,6 +5,7 @@ import logging
 import os
 import time
 import threading
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone, timedelta
 from decimal import Decimal, ROUND_DOWN
 from uuid import uuid4
@@ -22,9 +23,30 @@ FINISHED_ORDER_STATUS = "finished"
 _CANCEL_BREAKER_LOCK = threading.Lock()
 _CANCEL_BREAKER_NEXT_ALLOWED: dict[str, float] = {}
 _CANCEL_BREAKER_SECONDS = 3.0
+_SNAPSHOT_POOL = ThreadPoolExecutor(max_workers=1, thread_name_prefix="balance-snapshot")
 logger = logging.getLogger("exec_mcp")
 
 
+def _fire_balance_snapshot(account_id: str, source_ref: str) -> None:
+    """Schedule a balance snapshot in a background thread.
+
+    This avoids blocking the WebSocket event loop while the snapshot
+    capture (which makes HTTP calls and price lookups) runs.  The
+    single-worker pool serialises snapshots so they don't pile up.
+    """
+    try:
+        _SNAPSHOT_POOL.submit(_capture_snapshot_task, account_id, source_ref)
+    except RuntimeError:
+        pass  # pool already shut down
+
+
+def _capture_snapshot_task(account_id: str, source_ref: str) -> None:
+    try:
+        capture_account_balance_snapshot(account_id, source_kind="order_finished", source_ref=source_ref)
+    except Exception:
+        logger.debug("background balance snapshot failed for account_id=%s source_ref=%s", account_id, source_ref, exc_info=True)
+
+
 def _cancel_breaker_is_open(account_id: str) -> bool:
     with _CANCEL_BREAKER_LOCK:
         return _CANCEL_BREAKER_NEXT_ALLOWED.get(account_id, 0.0) > time.monotonic()
@@ -103,6 +125,12 @@ def _encode_raw_json(raw_json) -> str:
 
 
 def _record_order_status_update(*, bitstamp_order_id: str, status: str, raw_json, account_id: str | None = None) -> dict:
+    # ``account_id is None`` signals a WebSocket-driven call (the WS handler
+    # has no account context and looks it up from the order record below).
+    # Tool-driven calls (query_order, cancel_order, get_open_orders) always
+    # pass account_id explicitly.
+    from_ws = account_id is None
+
     record = _load_order_record(bitstamp_order_id=bitstamp_order_id)
     if record is None:
         raise HTTPException(status_code=404, detail="order not found")
@@ -121,10 +149,13 @@ def _record_order_status_update(*, bitstamp_order_id: str, status: str, raw_json
         conn.commit()
 
     if previous_status != FINISHED_ORDER_STATUS and normalized_status == FINISHED_ORDER_STATUS:
-        try:
-            capture_account_balance_snapshot(account_id, source_kind="order_finished", source_ref=bitstamp_order_id)
-        except Exception:
-            pass
+        if from_ws:
+            _fire_balance_snapshot(account_id, bitstamp_order_id)
+        else:
+            try:
+                capture_account_balance_snapshot(account_id, source_kind="order_finished", source_ref=bitstamp_order_id)
+            except Exception:
+                pass
 
     invalidate_account_cache(account_id)
     return {

+ 16 - 2
src/exec_mcp/storage.py

@@ -1,17 +1,31 @@
 from __future__ import annotations
 
 import sqlite3
+from contextlib import contextmanager
 from pathlib import Path
+from typing import Iterator
 
 DB_PATH = Path(__file__).resolve().parents[2] / "data" / "exec_mcp.sqlite3"
 
 
-def get_connection() -> sqlite3.Connection:
+@contextmanager
+def get_connection() -> Iterator[sqlite3.Connection]:
+    """Yield a SQLite connection that is always closed on exit.
+
+    The native ``sqlite3.Connection.__exit__`` only commits or rolls back —
+    it does **not** close the connection.  Every prior call site used
+    ``with get_connection() as conn:`` expecting full cleanup, so unclosed
+    connections accumulated over the process lifetime, each holding pager
+    cache memory.  This wrapper guarantees ``conn.close()`` runs on exit.
+    """
     DB_PATH.parent.mkdir(parents=True, exist_ok=True)
     conn = sqlite3.connect(DB_PATH)
     conn.row_factory = sqlite3.Row
     conn.execute("PRAGMA foreign_keys = ON")
-    return conn
+    try:
+        yield conn
+    finally:
+        conn.close()
 
 
 def _migrate_order_records(conn: sqlite3.Connection) -> None:

+ 54 - 0
tests/test_reporting_tools.py

@@ -249,6 +249,60 @@ def test_record_order_status_update_persists_raw_json_and_triggers_snapshot_once
     assert len(snapshot_calls) == 1
 
 
+def test_record_order_status_update_ws_path_defers_snapshot(temp_db, monkeypatch):
+    """When account_id is not passed (WS-driven call), the balance snapshot
+    is deferred to the background pool instead of running synchronously."""
+    account_id = _create_account()
+
+    with storage.get_connection() as conn:
+        conn.execute(
+            """
+            INSERT INTO order_records
+            (id, account_id, market, side, order_type, amount, price, expire_time, status, bitstamp_order_id, client_id, client_order_id, raw_json, created_at, updated_at)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                "rec-ws-1",
+                account_id,
+                "xrpusd",
+                "buy",
+                "limit",
+                "10",
+                "2.00",
+                None,
+                "open",
+                "1994292738961402",
+                "client-a",
+                None,
+                json.dumps({"status": "open"}),
+                "2026-04-09T08:20:50+00:00",
+                "2026-04-09T08:20:50+00:00",
+            ),
+        )
+        conn.commit()
+
+    submitted: list[tuple] = []
+
+    class _FakePool:
+        def submit(self, fn, *a):
+            submitted.append((fn, a))
+
+    monkeypatch.setattr(services_orders, "_SNAPSHOT_POOL", _FakePool())
+
+    # No account_id — simulates the WS handler call path
+    result = services_orders._record_order_status_update(
+        bitstamp_order_id="1994292738961402",
+        status="finished",
+        raw_json={"id": "1994292738961402", "status": "finished"},
+    )
+
+    assert result["ok"] is True
+    assert result["status"] == "finished"
+    assert len(submitted) == 1
+    assert submitted[0][0] is services_orders._capture_snapshot_task
+    assert submitted[0][1] == (account_id, "1994292738961402")
+
+
 def test_capture_account_balance_snapshot_uses_crypto_price(temp_db, monkeypatch):
     account_id = _create_account()