bug: `QuerySessionPool.acquire()` permanently leaks pool capacity when cancelled during session creation

[LuckySting]

TL;DR

QuerySessionPool.acquire() pre-increments _current_size before await self._create_new_session() and rolls it back only in an except Exception block. asyncio.CancelledError derives from BaseException, so a task cancelled while the session is being created leaves _current_size permanently incremented — one pool slot is lost forever. Once _current_size == size, every acquire() waits on _queue.get() indefinitely (there is no default acquire timeout): the pool is bricked until process restart.

• SDK version: reproduced on ydb==3.29.1, the code is unchanged on current main
• Mode: async (ydb.aio.query.QuerySessionPool)

Location

ydb-python-sdk/ydb/aio/query/pool.py

Lines 136 to 143 in 53b1b16

	self._current_size += 1
	try:
	session = await self._create_new_session()
	except Exception as e:
	# 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

self._current_size += 1
try:
    session = await self._create_new_session()
except Exception as e:
    # TODO: this exception could be retried via retrier, ...
    logger.warning("Failed to create new session")
    self._current_size -= 1
    raise e

How it happens in practice

In a grpc.aio service every client deadline/disconnect cancels the handler task, and the cancellation propagates into whatever the task is awaiting. If it lands inside _create_new_session() (a CreateSession RPC plus the AttachSession first response — a wide window on a loaded event loop), the accounting is silently corrupted.

The (correct) fix for #812 makes this path hotter: sessions invalidated by cancelled streams are returned to the queue dead, and each dead session popped by acquire() triggers a replacement create — precisely while more deadline-driven cancellations keep arriving. That produces a feedback loop: fewer real slots → acquires start queueing → more requests hit their deadline → more cancelled creates → more phantom slots. We drove a production-like pod (pool size 100, 50 concurrent read RPCs with a 4 s client deadline) to _current_size == 100 with an empty queue within a few minutes; it stayed there with zero traffic and every subsequent DB call hung forever.

Minimal reproduction (no network needed)

import asyncio
import ydb.aio.query.pool as pool_mod


async def slow_create(self):
    await asyncio.sleep(30)  # CreateSession + AttachSession in flight


pool_mod.QuerySessionPool._create_new_session = slow_create


async def main():
    pool = pool_mod.QuerySessionPool(None, size=1)

    task = asyncio.create_task(pool.acquire())
    await asyncio.sleep(0.05)  # let acquire() reach _create_new_session()
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        pass

    print(f"_current_size={pool._current_size} qsize={pool._queue.qsize()}")
    # -> _current_size=1 qsize=0 : pool believes it is full, yet no session exists

    await asyncio.wait_for(pool.acquire(), 1)  # -> TimeoutError: blocks forever


asyncio.run(main())

Output:

_current_size=1 qsize=0
TimeoutError

Proposed fix

Compensate the accounting for non-Exception BaseExceptions as well (keeping the current logging behavior for ordinary errors):

self._current_size += 1
try:
    session = await self._create_new_session()
except Exception as e:
    logger.warning("Failed to create new session")
    self._current_size -= 1
    raise e
except BaseException:  # asyncio.CancelledError and friends
    self._current_size -= 1
    raise

Related hardening in the same family: QuerySession.create() → _attach() also guards with except Exception only, so a cancellation between a successful CreateSession and attach completion additionally orphans the server-side session and the attach stream; the same BaseException treatment applies there.

Happy to send a PR with the fix and a regression test.

Related: #812 (dirty sessions returned to the pool on cancellation — fixed; this issue is the remaining accounting leak on the create path).