perf(selection): 优化选股执行性能

2026-08-12 09:45:16 +08:00
parent dd04933d63
commit 8963c067b3
22 changed files with 1333 additions and 120 deletions
@@ -0,0 +1,96 @@
+"""Bounded PostgreSQL pool ownership for selection adapters."""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any, Protocol, cast
+
+from psycopg_pool import ConnectionPool
+
+
+class SelectionConnectionPool(Protocol):
+    """Small pool surface shared by psycopg and unit-test fakes."""
+
+    def open(self, *, wait: bool = True) -> None: ...
+
+    def close(self) -> None: ...
+
+    def connection(self) -> Any: ...
+
+
+class SelectionPostgresPool:
+    """Own one bounded PostgreSQL pool for the selection read/write adapters.
+
+    The owner opens lazily on the first borrowed connection, which keeps app
+    construction cheap while still making the pool lifetime process-scoped in
+    the HTTP composition layer.  A fake pool can be injected for unit tests.
+    """
+
+    def __init__(
+        self,
+        database_url: str,
+        *,
+        max_connections: int,
+        pool: SelectionConnectionPool | None = None,
+    ) -> None:
+        """Create a bounded pool owner with an optional injected pool."""
+
+        if max_connections < 1:
+            raise ValueError("max_connections must be at least 1")
+        self.database_url = database_url
+        self.max_connections = max_connections
+        self.pool: SelectionConnectionPool = pool or cast(
+            SelectionConnectionPool,
+            ConnectionPool(
+                conninfo=database_url,
+                min_size=1,
+                max_size=max_connections,
+                open=False,
+            ),
+        )
+        self._pool_open = False
+        self._pool_state_lock = threading.Lock()
+
+    def open(self) -> None:
+        """Open the underlying pool once and wait for its minimum connection."""
+
+        with self._pool_state_lock:
+            if self._pool_open:
+                return
+            if bool(getattr(self.pool, "_opened", False)):
+                self._pool_open = True
+                return
+            self.pool.open(wait=True)
+            self._pool_open = True
+
+    def close(self) -> None:
+        """Close the pool after borrowed connections have been returned."""
+
+        with self._pool_state_lock:
+            if self._pool_open or bool(getattr(self.pool, "_opened", False)):
+                self.pool.close()
+                self._pool_open = False
+
+    @contextmanager
+    def connection(self) -> Generator[Any, None, None]:
+        """Borrow one connection and return it to the bounded pool."""
+
+        self.open()
+        with self.pool.connection() as connection:
+            yield connection
+
+    def __enter__(self) -> SelectionPostgresPool:
+        """Open and return this resource owner."""
+
+        self.open()
+        return self
+
+    def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None:
+        """Release the owned pool at the end of a context."""
+
+        self.close()
+
+
+__all__ = ["SelectionConnectionPool", "SelectionPostgresPool"]