Do not leak query pool capacity on cancelled session create

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed async `QuerySessionPool` permanently losing a pool slot when `acquire()` was cancelled while a new session was being created: `asyncio.CancelledError` no longer leaks the pool size counter, so a pool under deadline-driven cancellations can no longer end up exhausted and blocking forever. A cancelled or interrupted session attach now also closes the session instead of orphaning it server-side
+
 ## 3.31.2 ##
 * Add `TableClient.describe_system_view` (sync and async) returning a dedicated `SystemViewSchemeEntry` with the system view id, name, columns, primary key and attributes — a full description of system view objects that `describe_table` does not provide
 * Add `pool_id` parameter to `QuerySession.execute`, `QueryTxContext.execute`, and `QuerySessionPool.execute_with_retries` to route queries to a specific resource pool

ydb/aio/query/pool.py

@@ -136,11 +136,17 @@ async def acquire(self, timeout: Optional[float] = None) -> QuerySession:
         self._current_size += 1
         try:
             session = await self._create_new_session()
-        except Exception as e:
+        except Exception:
             # TODO: this exception could be retried via retrier, so no need to log error here. Probably we should retry this right in create_new_session method.
             logger.warning("Failed to create new session")
             self._current_size -= 1
-            raise e
+            raise
+        except BaseException:
+            # asyncio.CancelledError does not derive from Exception, so without this
+            # branch a task cancelled while its session is being created would leave
+            # _current_size incremented forever and permanently lose a pool slot.
+            self._current_size -= 1
+            raise
 
         return session
 

ydb/aio/query/pool_test.py

@@ -5,6 +5,7 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 from ydb import issues
+from ydb.aio import _utilities as aio_utilities
 from ydb.aio.query.pool import QuerySessionPool
 from ydb.aio.query.session import QuerySession
 from ydb.aio.query.transaction import QueryTxContext
@@ -32,6 +33,11 @@ def _make_active_session():
     return session
 
 
+async def _wait_signalled(event: asyncio.Event) -> None:
+    """Wait for a test signal, bounded so a regression fails the test instead of hanging the suite."""
+    await asyncio.wait_for(event.wait(), timeout=5)
+
+
 class TestAcquireTimeout(unittest.IsolatedAsyncioTestCase):
     async def test_acquire_returns_session_when_available(self):
         pool = _make_pool(size=2)
@@ -121,6 +127,111 @@ async def test_retry_reacquires_invalidated_session_before_first_use(self):
         live_session.explain.assert_awaited_once_with("SELECT 1")
 
 
+class TestAcquireCancellation(unittest.IsolatedAsyncioTestCase):
+    """Cancelling acquire() while the session is being created must not lose a pool slot."""
+
+    @staticmethod
+    def _hanging_create(entered: asyncio.Event):
+        async def slow_create():
+            entered.set()
+            await asyncio.sleep(30)
+
+        return slow_create
+
+    async def _cancel_during_create(self, pool):
+        entered = asyncio.Event()
+        pool._create_new_session = self._hanging_create(entered)
+
+        task = asyncio.create_task(pool.acquire())
+        await _wait_signalled(entered)
+        task.cancel()
+        with self.assertRaises(asyncio.CancelledError):
+            await task
+
+    async def test_cancelled_create_does_not_leak_pool_capacity(self):
+        pool = _make_pool(size=1)
+
+        await self._cancel_during_create(pool)
+
+        self.assertEqual(pool._current_size, 0)
+
+    async def test_pool_still_usable_after_cancelled_create(self):
+        pool = _make_pool(size=1)
+
+        await self._cancel_during_create(pool)
+
+        session = _make_active_session()
+        pool._create_new_session = AsyncMock(return_value=session)
+
+        acquired = await asyncio.wait_for(pool.acquire(), timeout=1)
+
+        self.assertIs(acquired, session)
+        self.assertEqual(pool._current_size, 1)
+
+    async def test_failed_create_does_not_leak_pool_capacity(self):
+        pool = _make_pool(size=1)
+        pool._create_new_session = AsyncMock(side_effect=issues.ConnectionError("no connection"))
+
+        with self.assertRaises(issues.ConnectionError):
+            await pool.acquire()
+
+        self.assertEqual(pool._current_size, 0)
+
+
+class TestSessionAttachCancellation(unittest.IsolatedAsyncioTestCase):
+    """A cancelled attach must retire the session instead of orphaning it server-side."""
+
+    def _make_session(self):
+        driver = MagicMock()
+        driver._driver_config.query_client_settings = None
+        session = QuerySession(driver)
+        session._session_id = "fake-session-id"
+        return session
+
+    async def test_attach_invalidates_session_when_cancelled_awaiting_first_response(self):
+        session = self._make_session()
+        stream = MagicMock()
+        entered = asyncio.Event()
+
+        async def fake_attach_call(*args, **kwargs):
+            return stream
+
+        async def hanging_first_message(*args, **kwargs):
+            entered.set()
+            await asyncio.sleep(30)
+
+        with patch.object(type(session), "_attach_call", side_effect=fake_attach_call), patch.object(
+            aio_utilities, "AsyncResponseIterator", MagicMock()
+        ), patch.object(aio_utilities, "get_first_message_with_timeout", hanging_first_message):
+            task = asyncio.create_task(session._attach())
+            await _wait_signalled(entered)
+            task.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await task
+
+        self.assertFalse(session.is_active)
+        self.assertTrue(session._invalidated)
+        stream.cancel.assert_called_once()
+
+    async def test_attach_invalidates_session_when_cancelled_before_stream_is_open(self):
+        session = self._make_session()
+        entered = asyncio.Event()
+
+        async def hanging_attach_call(*args, **kwargs):
+            entered.set()
+            await asyncio.sleep(30)
+
+        with patch.object(type(session), "_attach_call", side_effect=hanging_attach_call):
+            task = asyncio.create_task(session._attach())
+            await _wait_signalled(entered)
+            task.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await task
+
+        self.assertFalse(session.is_active)
+        self.assertTrue(session._invalidated)
+
+
 async def _async_empty_iter():
     """Async-iterable that yields nothing; usable as a stub for session.execute return value."""
     if False:

ydb/aio/query/session.py

@@ -49,21 +49,23 @@ def __init__(
         self._status_stream = None
 
     async def _attach(self) -> None:
-        self._stream = await self._attach_call()
-        self._status_stream = _utilities.AsyncResponseIterator(
-            self._stream,
-            self._attach_stream_wrapper,
-        )
-
         try:
+            self._stream = await self._attach_call()
+            self._status_stream = _utilities.AsyncResponseIterator(
+                self._stream,
+                self._attach_stream_wrapper,
+            )
+
             first_response = await _utilities.get_first_message_with_timeout(
                 self._status_stream,
                 DEFAULT_INITIAL_RESPONSE_TIMEOUT,
             )
             issues._process_response(first_response)
-        except Exception as e:
+        except BaseException:
+            # BaseException, not Exception: a cancelled attach must tear the stream
+            # down too, otherwise the half-attached session is orphaned server-side.
             self._close_session(invalidate=True)
-            raise e
+            raise
 
         self._loop.create_task(self._check_session_status_loop(), name="check session status task")
 

ydb/query/pool_test.py

@@ -8,7 +8,7 @@
 
 from unittest.mock import patch
 
-from ydb import issues
+from ydb import _utilities, issues
 from ydb.convert import _ResultSet, aggregate_result_sets_by_index, aggregate_result_sets_by_index_async
 from ydb.query.base import create_execute_query_request
 from ydb.query.pool import QuerySessionPool
@@ -64,6 +64,41 @@ def release_after_delay():
             t.join()
 
 
+class TestSessionAttachInterrupted(unittest.TestCase):
+    """An interrupted attach must retire the session instead of orphaning it server-side."""
+
+    def _make_session(self):
+        driver = MagicMock()
+        driver._driver_config.query_client_settings = None
+        session = QuerySession(driver)
+        session._session_id = "fake-session-id"
+        return session
+
+    def test_attach_invalidates_session_when_interrupted_awaiting_first_response(self):
+        session = self._make_session()
+        stream = MagicMock()
+
+        with patch.object(type(session), "_attach_call", return_value=stream), patch.object(
+            _utilities, "SyncResponseIterator", MagicMock()
+        ), patch.object(_utilities, "get_first_message_with_timeout", side_effect=KeyboardInterrupt):
+            with self.assertRaises(KeyboardInterrupt):
+                session._attach()
+
+        self.assertFalse(session.is_active)
+        self.assertTrue(session._invalidated)
+        stream.cancel.assert_called_once()
+
+    def test_attach_invalidates_session_when_interrupted_before_stream_is_open(self):
+        session = self._make_session()
+
+        with patch.object(type(session), "_attach_call", side_effect=KeyboardInterrupt):
+            with self.assertRaises(KeyboardInterrupt):
+                session._attach()
+
+        self.assertFalse(session.is_active)
+        self.assertTrue(session._invalidated)
+
+
 def _rs(index, rows, columns=None, truncated=False, data=None):
     return _ResultSet(
         columns=["id"] if columns is None else columns,

ydb/query/session.py

@@ -385,21 +385,23 @@ def __init__(self, driver: "SyncDriver", settings: Optional[base.QueryClientSett
         super().__init__(driver, settings)
 
     def _attach(self, first_resp_timeout: int = DEFAULT_INITIAL_RESPONSE_TIMEOUT) -> None:
-        self._stream = self._attach_call()
-        status_stream = _utilities.SyncResponseIterator(
-            self._stream,
-            self._attach_stream_wrapper,
-        )
-
         try:
+            self._stream = self._attach_call()
+            status_stream = _utilities.SyncResponseIterator(
+                self._stream,
+                self._attach_stream_wrapper,
+            )
+
             first_response = _utilities.get_first_message_with_timeout(
                 status_stream,
                 first_resp_timeout,
             )
             issues._process_response(first_response)
-        except Exception as e:
+        except BaseException:
+            # BaseException, not Exception: an interrupted attach must tear the stream
+            # down too, otherwise the half-attached session is orphaned server-side.
             self._close_session(invalidate=True)
-            raise e
+            raise
 
         threading.Thread(
             target=self._check_session_status_loop,