From ede160617dbc26989a9ebbc170b4f684909e4c19 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 6 Aug 2026 14:53:13 +0300 Subject: [PATCH 1/2] Do not leak query pool capacity on cancelled session create asyncio.CancelledError is not an Exception, so cancelling acquire() while _create_new_session() was in flight left _current_size incremented forever and permanently lost a pool slot. Also close the session when attach is cancelled or interrupted, instead of orphaning it server-side. Fixes #870 --- CHANGELOG.md | 2 + ydb/aio/query/pool.py | 6 +++ ydb/aio/query/pool_test.py | 106 +++++++++++++++++++++++++++++++++++++ ydb/aio/query/session.py | 18 ++++--- ydb/query/pool_test.py | 37 ++++++++++++- ydb/query/session.py | 18 ++++--- 6 files changed, 170 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 39e6077d..c98d4276 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/ydb/aio/query/pool.py b/ydb/aio/query/pool.py index 25f92703..48956036 100644 --- a/ydb/aio/query/pool.py +++ b/ydb/aio/query/pool.py @@ -141,6 +141,12 @@ async def acquire(self, timeout: Optional[float] = None) -> QuerySession: logger.warning("Failed to create new session") self._current_size -= 1 raise e + 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 diff --git a/ydb/aio/query/pool_test.py b/ydb/aio/query/pool_test.py index 71d083da..4c873f37 100644 --- a/ydb/aio/query/pool_test.py +++ b/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 @@ -121,6 +122,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 entered.wait() + 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 entered.wait() + 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 entered.wait() + 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: diff --git a/ydb/aio/query/session.py b/ydb/aio/query/session.py index 7f765458..c1a1fea7 100644 --- a/ydb/aio/query/session.py +++ b/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") diff --git a/ydb/query/pool_test.py b/ydb/query/pool_test.py index 33041ccf..599b0896 100644 --- a/ydb/query/pool_test.py +++ b/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, diff --git a/ydb/query/session.py b/ydb/query/session.py index b8a7d82e..1cec3600 100644 --- a/ydb/query/session.py +++ b/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, From 8a4fd0f9531af1cecfcb61e1f3a9628ae5ca67c1 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 6 Aug 2026 15:30:54 +0300 Subject: [PATCH 2/2] Address review: bare raise, bound test waits --- ydb/aio/query/pool.py | 4 ++-- ydb/aio/query/pool_test.py | 11 ++++++++--- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/ydb/aio/query/pool.py b/ydb/aio/query/pool.py index 48956036..a52c689f 100644 --- a/ydb/aio/query/pool.py +++ b/ydb/aio/query/pool.py @@ -136,11 +136,11 @@ 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 diff --git a/ydb/aio/query/pool_test.py b/ydb/aio/query/pool_test.py index 4c873f37..e30a2c93 100644 --- a/ydb/aio/query/pool_test.py +++ b/ydb/aio/query/pool_test.py @@ -33,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) @@ -138,7 +143,7 @@ async def _cancel_during_create(self, pool): pool._create_new_session = self._hanging_create(entered) task = asyncio.create_task(pool.acquire()) - await entered.wait() + await _wait_signalled(entered) task.cancel() with self.assertRaises(asyncio.CancelledError): await task @@ -199,7 +204,7 @@ async def hanging_first_message(*args, **kwargs): aio_utilities, "AsyncResponseIterator", MagicMock() ), patch.object(aio_utilities, "get_first_message_with_timeout", hanging_first_message): task = asyncio.create_task(session._attach()) - await entered.wait() + await _wait_signalled(entered) task.cancel() with self.assertRaises(asyncio.CancelledError): await task @@ -218,7 +223,7 @@ async def hanging_attach_call(*args, **kwargs): with patch.object(type(session), "_attach_call", side_effect=hanging_attach_call): task = asyncio.create_task(session._attach()) - await entered.wait() + await _wait_signalled(entered) task.cancel() with self.assertRaises(asyncio.CancelledError): await task