Move node pessimation to connection pools

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: KirillKurdyukov

tests/query/test_query_session_hints.py

@@ -1,14 +1,17 @@
+import asyncio
 from unittest import mock
 
+import pytest
+
 from ydb._grpc.common.protos import ydb_query_pb2
+from ydb.aio.pool import ConnectionPool as AsyncConnectionPool
+from ydb.pool import ConnectionPool
 from ydb.query.session import QuerySession
 
 
 def _make_session(node_id=42):
     driver = mock.Mock()
-    driver._store = mock.Mock()
-    driver._store.connections_by_node_id = {node_id: mock.Mock()}
-    driver._on_disconnected = mock.Mock(return_value=None)
+    driver._pessimize_node = mock.Mock()
     session = QuerySession(driver)
     session._session_id = "test-session"
     session._node_id = node_id
@@ -18,7 +21,6 @@ def _make_session(node_id=42):
 class TestQuerySessionAttachHints:
     def test_node_shutdown_pessimizes_node_and_invalidates_session(self):
         session, driver = _make_session(node_id=42)
-        connection = driver._store.connections_by_node_id[42]
 
         session._handle_attach_session_state(
             ydb_query_pb2.SessionState(
@@ -27,7 +29,7 @@ def test_node_shutdown_pessimizes_node_and_invalidates_session(self):
             )
         )
 
-        driver._on_disconnected.assert_called_once_with(connection)
+        driver._pessimize_node.assert_called_once_with(42)
         assert session._invalidated
         assert session._closed
 
@@ -41,11 +43,11 @@ def test_session_shutdown_invalidates_without_pessimizing_node(self):
             )
         )
 
-        driver._on_disconnected.assert_not_called()
+        driver._pessimize_node.assert_not_called()
         assert session._invalidated
         assert session._closed
 
-    def test_node_shutdown_with_zero_node_id_skips_pessimization(self):
+    def test_node_shutdown_with_zero_node_id_delegates_to_driver(self):
         session, driver = _make_session(node_id=0)
 
         session._handle_attach_session_state(
@@ -55,7 +57,7 @@ def test_node_shutdown_with_zero_node_id_skips_pessimization(self):
             )
         )
 
-        driver._on_disconnected.assert_not_called()
+        driver._pessimize_node.assert_called_once_with(0)
         assert session._invalidated
 
     def test_no_hint_does_not_invalidate(self):
@@ -65,6 +67,46 @@ def test_no_hint_does_not_invalidate(self):
             ydb_query_pb2.SessionState(status=0),
         )
 
-        driver._on_disconnected.assert_not_called()
+        driver._pessimize_node.assert_not_called()
         assert not session._invalidated
         assert not session._closed
+
+
+class TestConnectionPoolAttachHintPessimization:
+    def test_sync_pool_pessimizes_node_connection(self):
+        pool = ConnectionPool.__new__(ConnectionPool)
+        connection = mock.Mock()
+        pool._store = mock.Mock()
+        pool._store.connections_by_node_id = {42: connection}
+        pool._on_disconnected = mock.Mock()
+
+        pool._pessimize_node(42)
+
+        pool._on_disconnected.assert_called_once_with(connection)
+
+    def test_sync_pool_ignores_missing_node_connection(self):
+        pool = ConnectionPool.__new__(ConnectionPool)
+        pool._store = mock.Mock()
+        pool._store.connections_by_node_id = {}
+        pool._on_disconnected = mock.Mock()
+
+        pool._pessimize_node(42)
+        pool._pessimize_node(0)
+        pool._pessimize_node(None)
+
+        pool._on_disconnected.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_async_pool_pessimizes_node_connection(self):
+        pool = AsyncConnectionPool.__new__(AsyncConnectionPool)
+        connection = mock.Mock()
+        disconnect = mock.AsyncMock()
+        pool._store = mock.Mock()
+        pool._store.connections_by_node_id = {42: connection}
+        pool._on_disconnected = mock.Mock(return_value=disconnect)
+
+        pool._pessimize_node(42)
+        await asyncio.sleep(0)
+
+        pool._on_disconnected.assert_called_once_with(connection)
+        disconnect.assert_awaited_once_with()

ydb/aio/pool.py

@@ -3,7 +3,7 @@
 import asyncio
 import logging
 import random
-from typing import Any, Callable, Optional, Tuple, TYPE_CHECKING
+from typing import Any, Callable, Optional, Tuple, TYPE_CHECKING, cast
 
 from ydb import issues
 from ydb.opentelemetry.tracing import SpanName, create_ydb_span
@@ -311,6 +311,15 @@ async def __wrapper__() -> None:
 
         return __wrapper__
 
+    def _pessimize_node(self, node_id: Optional[int]) -> None:
+        """Deprioritize the connection attached to the given YDB node."""
+        if node_id is None or node_id <= 0:
+            return
+
+        connection = cast(Optional[Connection], self._store.connections_by_node_id.get(node_id))
+        if connection is not None:
+            asyncio.get_running_loop().create_task(self._on_disconnected(connection)())
+
     async def wait(self, timeout: Optional[float] = 7.0, fail_fast: bool = False) -> None:  # type: ignore[override]  # async override of sync method
         with create_ydb_span(SpanName.DRIVER_INITIALIZE, self._driver_config, kind="internal").attach_context():
             await self._store.get(fast_fail=fail_fast, wait_timeout=timeout if timeout is not None else 7.0)

ydb/pool.py

@@ -494,6 +494,15 @@ def _on_disconnected(self, connection: Connection) -> None:
         if self._discovery_thread:
             self._discovery_thread.notify_disconnected()
 
+    def _pessimize_node(self, node_id: Optional[int]) -> None:
+        """Deprioritize the connection attached to the given YDB node."""
+        if node_id is None or node_id <= 0:
+            return
+
+        connection = self._store.connections_by_node_id.get(node_id)
+        if connection is not None:
+            self._on_disconnected(connection)
+
     def discovery_debug_details(self) -> str:
         """
         Returns debug string about last errors

ydb/query/session.py

@@ -1,5 +1,4 @@
 import abc
-import asyncio
 import json
 import logging
 import threading
@@ -169,42 +168,11 @@ def _handle_attach_session_state(self, response_pb) -> None:
 
         hint = response_pb.WhichOneof("session_hint")
         if hint == "node_shutdown":
-            self._pessimize_node_by_id(self._node_id)
+            self._driver._pessimize_node(self._node_id)
             self._close_session(invalidate=True)
         elif hint == "session_shutdown":
             self._close_session(invalidate=True)
 
-    def _pessimize_node_by_id(self, node_id: Optional[int]) -> None:
-        """Deprioritize the node's connection (same effect as a transport failure)."""
-        if node_id is None or node_id <= 0:
-            return
-
-        store = getattr(self._driver, "_store", None)
-        if store is None:
-            return
-
-        connection = getattr(store, "connections_by_node_id", {}).get(node_id)
-        if connection is None:
-            return
-
-        on_disconnected = getattr(self._driver, "_on_disconnected", None)
-        if on_disconnected is None:
-            connection.close()
-            return
-
-        disconnect_cb = on_disconnected(connection)
-        if disconnect_cb is None:
-            return
-
-        if asyncio.iscoroutinefunction(disconnect_cb):
-            loop = getattr(self, "_loop", None)
-            if loop is not None and loop.is_running():
-                loop.create_task(disconnect_cb())
-            return
-
-        if callable(disconnect_cb):
-            disconnect_cb()
-
     def _close_session(self, invalidate: bool = False) -> None:
         if self._closed:
             return