|
29 | 29 |
|
30 | 30 | __all__ = ["PostgresBlockStore", "BlockStoreError"] |
31 | 31 |
|
| 32 | +# ─── Process-wide connection-pool registry (thread-leak fix) ───────────────── |
| 33 | +# |
| 34 | +# ``storage.get_block_store()`` is a factory: every recall() / hybrid_search() |
| 35 | +# / _check_workspace() call on a Postgres-backed workspace constructs a |
| 36 | +# *fresh* PostgresBlockStore (by design — construction itself is free, see |
| 37 | +# __init__ below). ``_get_pool()`` used to open a brand-new |
| 38 | +# ``psycopg_pool.ConnectionPool(..., open=True)`` on that fresh instance's |
| 39 | +# first real query. Opening a pool spawns a scheduler thread plus |
| 40 | +# ``num_workers`` (default 3) worker threads that run until ``.close()`` is |
| 41 | +# called — and nothing on the per-MCP-tool-call path ever called it, since |
| 42 | +# the ephemeral PostgresBlockStore went out of scope as soon as the tool |
| 43 | +# call returned. The 4 background threads per call were only reclaimed |
| 44 | +# incidentally whenever Python's garbage collector happened to collect the |
| 45 | +# discarded pool, which is not guaranteed to keep pace under sustained |
| 46 | +# traffic. Production symptom: 76,479 accumulated threads / ~32GB RSS over |
| 47 | +# 2.6 days on a long-running `mcp_server.py` process (~1 leaked thread set |
| 48 | +# per tool call), exhausting the box's fork/thread capacity. |
| 49 | +# |
| 50 | +# Fix: key pools by (dsn, schema) in a process-wide registry so every |
| 51 | +# PostgresBlockStore pointed at the same database shares one long-lived |
| 52 | +# pool, regardless of how many wrapper instances the factory constructs. |
| 53 | +_pool_registry: dict[tuple[str, str], Any] = {} |
| 54 | +_pool_registry_lock = threading.Lock() |
| 55 | + |
32 | 56 | # Schema names must be safe Postgres identifiers (no injection surface). |
33 | 57 | _SAFE_SCHEMA_RE = re.compile(r"^[a-z_][a-z0-9_]{0,62}$") |
34 | 58 |
|
@@ -531,25 +555,42 @@ def __init__( |
531 | 555 | # ─── Lifecycle ──────────────────────────────────────────────────────────── |
532 | 556 |
|
533 | 557 | def _get_pool(self) -> Any: |
534 | | - """Return the connection pool, creating it on first call.""" |
| 558 | + """Return the connection pool, creating (or reusing) it on first call. |
| 559 | +
|
| 560 | + Pools are cached process-wide in ``_pool_registry``, keyed by |
| 561 | + ``(dsn, schema)`` — see the module-level note above. This avoids |
| 562 | + opening a brand-new ``ConnectionPool`` (and its background |
| 563 | + scheduler + worker threads) for every ephemeral PostgresBlockStore |
| 564 | + the ``storage.get_block_store()`` factory constructs. ``self._pool`` |
| 565 | + still caches the resolved pool locally so repeated calls on the |
| 566 | + same instance skip the registry lookup. |
| 567 | + """ |
535 | 568 | if self._pool is not None: |
536 | 569 | return self._pool |
537 | 570 | _, ConnectionPool = _require_psycopg() |
538 | 571 | with self._init_lock: |
539 | | - if self._pool is None: |
540 | | - self._pool = ConnectionPool( |
541 | | - self._dsn, |
542 | | - min_size=1, |
543 | | - max_size=10, |
544 | | - open=True, |
545 | | - # Append the pgvector extension's schema to each |
546 | | - # connection's search_path so the bare ``vector`` type, |
547 | | - # ``<=>`` operator and ``vector_cosine_ops`` opclass |
548 | | - # resolve even when the DSN isolates search_path to a |
549 | | - # single workspace schema (bug-6: embeddings silently |
550 | | - # disabled / hybrid_search crashed under isolation). |
551 | | - configure=_configure_vector_search_path, |
552 | | - ) |
| 572 | + if self._pool is not None: |
| 573 | + return self._pool |
| 574 | + key = (self._dsn, self._schema) |
| 575 | + with _pool_registry_lock: |
| 576 | + pool = _pool_registry.get(key) |
| 577 | + if pool is None or pool.closed: |
| 578 | + pool = ConnectionPool( |
| 579 | + self._dsn, |
| 580 | + min_size=1, |
| 581 | + max_size=10, |
| 582 | + open=True, |
| 583 | + # Append the pgvector extension's schema to each |
| 584 | + # connection's search_path so the bare ``vector`` |
| 585 | + # type, ``<=>`` operator and ``vector_cosine_ops`` |
| 586 | + # opclass resolve even when the DSN isolates |
| 587 | + # search_path to a single workspace schema (bug-6: |
| 588 | + # embeddings silently disabled / hybrid_search |
| 589 | + # crashed under isolation). |
| 590 | + configure=_configure_vector_search_path, |
| 591 | + ) |
| 592 | + _pool_registry[key] = pool |
| 593 | + self._pool = pool |
553 | 594 | return self._pool |
554 | 595 |
|
555 | 596 | def ping(self, *, timeout: float = 5.0) -> dict[str, Any]: |
@@ -1261,12 +1302,24 @@ def list_files(self) -> list[str]: |
1261 | 1302 | # ─── Context manager / cleanup ──────────────────────────────────────────── |
1262 | 1303 |
|
1263 | 1304 | def close(self) -> None: |
1264 | | - """Close the underlying connection pool.""" |
| 1305 | + """Close the underlying connection pool. |
| 1306 | +
|
| 1307 | + Pools are shared process-wide via ``_pool_registry`` (keyed by |
| 1308 | + ``(dsn, schema)``), so this also evicts the registry entry — |
| 1309 | + otherwise a *different* PostgresBlockStore instance still holding |
| 1310 | + the same (now-closed) pool object cached on ``self._pool`` would |
| 1311 | + try to reuse it. The next ``_get_pool()`` call for this |
| 1312 | + ``(dsn, schema)`` opens a fresh pool. |
| 1313 | + """ |
1265 | 1314 | if self._pool is not None: |
| 1315 | + key = (self._dsn, self._schema) |
1266 | 1316 | try: |
1267 | 1317 | self._pool.close() |
1268 | 1318 | except Exception as exc: |
1269 | 1319 | _log.debug("pg_pool_close_failed: %s", exc) |
| 1320 | + with _pool_registry_lock: |
| 1321 | + if _pool_registry.get(key) is self._pool: |
| 1322 | + del _pool_registry[key] |
1270 | 1323 | self._pool = None |
1271 | 1324 |
|
1272 | 1325 | def __enter__(self) -> "PostgresBlockStore": |
|
0 commit comments