bug: `asyncio.CancelledError` in `QuerySessionPool` leaves sessions dirty

[vladkolotvin]

asyncio.CancelledError in QuerySessionPool leaves sessions dirty → SessionBusy: Pending previous query completion

TL;DR

When a task inside QuerySessionPool.retry_tx_async is cancelled during tx.execute(...) (via task.cancel()), the underlying gRPC stream can be left undrained, the tx can be left un-rolled-back, the session is put back into the pool as-is, and the next consumer of that session fails with:

ydb.issues.SessionBusy: message: "Pending previous query completion" severity: 1 (server_code: 400190)

At that point the problem is already inside the SDK state: the pool still holds a session that the server considers busy with the previous request.

This report documents five independent deficiencies in the async Query pool code path that together produce this failure mode. Each case below is either observed directly in the e2e reproduction or tied to a minimal, verified reproduction against the unmodified SDK.

• SDK version: ydb==3.26.4
• Python: 3.12
• Mode: async (ydb.aio.query.QuerySessionPool)
• Verified against a real ydb docker image (YDB_ENDPOINT=grpc://ydb:2136, YDB_DATABASE=/local).

All outputs shown below are copied from an actual run — not hypothetical.

Real-world traceback

From a production service where the task running a YDB transaction gets cancelled:

File ".../ydb/aio/query/pool.py", line 181, in retry_tx_async
    return await retry_operation_async(wrapped_callee, retry_settings)
File ".../ydb/retries.py", line 195, in retry_operation_async
    return await next_opt.result
File ".../ydb/aio/query/pool.py", line 178, in wrapped_callee
    await tx.commit()
File ".../ydb/aio/query/transaction.py", line 111, in commit
    await self._commit_call(settings)
...
ydb.issues.SessionBusy: message: "Pending previous query completion" severity: 1 (server_code: 400190)

The immediate caller that fails is tx.commit() on a freshly checked-out session. The client has not issued anything yet on this session in this attempt — the "previous query" belongs to a previous, cancelled, completely unrelated request that ran on the same session. A single mis-timed cancel poisons a pool slot for multiple subsequent unrelated callers.

End-to-end reproduction

A single surgical monkey-patch simulates what grpc.aio actually does when the underlying call is cancelled: every subsequent read on the stream raises CancelledError until the call is torn down. Everything else below is stock SDK + a real YDB docker — no task.cancel(), no fake iterators, no custom pool.

One such cancellation inside a single tx.execute iteration is enough to:

• bypass session._on_execute_stream_error (case 1),
• leave the server-side stream undrained (case 3),
• leave the transaction open with no rollback (case 4),
• drop the poisoned session straight back into the pool (case 5),
• and poison the next caller that lands on that pool slot.

# repro_e2e.py  —  ydb==3.26.4
import asyncio
import os

from ydb.aio import Driver, QuerySessionPool
from ydb.aio._utilities import AsyncResponseIterator


armed = False
_orig_next = AsyncResponseIterator._next


async def patched_next(self):
    if armed:
        # Once a grpc.aio call is cancelled, every subsequent read on it
        # raises CancelledError until the call is fully torn down.
        raise asyncio.CancelledError('simulated grpc.aio call cancellation')
    return await _orig_next(self)


AsyncResponseIterator._next = patched_next


async def main():
    global armed
    async with Driver(endpoint=os.environ['YDB_ENDPOINT'], database=os.environ['YDB_DATABASE']) as driver:
        await driver.wait(5, fail_fast=True)
        async with QuerySessionPool(driver, size=1) as pool:
            victim_sid = None
            on_error_calls = []
            rollback_calls = []

            async def victim(tx):
                nonlocal victim_sid
                victim_sid = tx.session.session_id

                orig_hook = tx.session._on_execute_stream_error
                def spy_hook(e):
                    on_error_calls.append(type(e).__name__)
                    return orig_hook(e)
                tx.session._on_execute_stream_error = spy_hook

                orig_rollback = tx.rollback
                async def spy_rollback(*a, **kw):
                    rollback_calls.append(True)
                    return await orig_rollback(*a, **kw)
                tx.rollback = spy_rollback

                global armed
                result = await tx.execute('SELECT 1')
                armed = True  # stays armed through tx/stream teardown
                async for _ in result:
                    pass

            try:
                await pool.retry_tx_async(victim)
            except BaseException:
                pass
            armed = False  # stop simulating cancellation for the probe

            print('--- observations on the cancelled call ---')
            print(f'session._on_execute_stream_error invoked: {on_error_calls}')
            print(f'tx.rollback() called:                     {bool(rollback_calls)}')

            qsz = pool._queue.qsize()
            sess_in_queue = list(pool._queue._queue)[0] if qsz else None
            print()
            print('--- pool state after cancellation ---')
            print(f'pool queue size:        {qsz}')
            print(f'same session re-queued: {sess_in_queue is not None and sess_in_queue.session_id == victim_sid}')
            print(f'session.is_active:      {sess_in_queue.is_active if sess_in_queue else None}')

            async with pool.checkout() as sess:
                reused = sess.session_id == victim_sid
                tx = sess.transaction()
                probe_err = None
                try:
                    await tx.begin()
                    await tx.execute('SELECT 1')
                except BaseException as e:
                    probe_err = f'{type(e).__name__}: {str(e).splitlines()[0]}'
                print()
                print('--- next unrelated caller ---')
                print(f'next checkout == victim: {reused}')
                print(f'first RPC:               {probe_err or "OK"}')


if __name__ == '__main__':
    asyncio.run(main())

Actual output:

--- observations on the cancelled call ---
session._on_execute_stream_error invoked: []
tx.rollback() called:                     False

--- pool state after cancellation ---
pool queue size:        1
same session re-queued: True
session.is_active:      True

--- next unrelated caller ---
next checkout == victim: True
first RPC:               SessionBusy: message: "Pending previous query completion" severity: 1 (server_code: 400190)

Reading the output top to bottom:

• session._on_execute_stream_error invoked: [] — case 1: the stream hook was never called, because AsyncResponseContextIterator._next catches Exception, not BaseException.
• tx.rollback() called: False — case 4: the transaction was never rolled back, because QueryTxContext.__aexit__ re-raises the drain's CancelledError before the rollback block.
• same session re-queued: True + session.is_active: True — case 5: SimpleQuerySessionCheckoutAsync.__aexit__ dropped the poisoned session back into the FIFO queue with no sanity check; the SDK believes it is healthy.
• first RPC: SessionBusy: Pending previous query completion — case 3: the server-side stream was never fully drained, so the next caller lands on a session that the server still considers busy with the previous query.

The sections below annotate each contributing deficiency with the relevant SDK code. Cases 1, 3, 4, 5 are all observed in the e2e output above; case 2 is an independent amplifier with its own e2e reproduction.

---

Case 1. AsyncResponseContextIterator._next catches Exception, so CancelledError bypasses on_error

ydb/aio/query/base.py:

class AsyncResponseContextIterator(_utilities.AsyncResponseIterator):
    ...
    async def _next(self):
        try:
            return await super()._next()
        except Exception as e:
            if self._on_error:
                self._on_error(e)
            raise e

asyncio.CancelledError inherits from BaseException, not Exception, so it bypasses this handler entirely. The _on_error hook — which is bound to session._on_execute_stream_error — is never invoked when an in-flight stream is interrupted by a cancel. This is exactly what the e2e run above shows: session._on_execute_stream_error invoked: [].

---

Case 2. BaseQuerySession._on_execute_stream_error only invalidates on DeadlineExceed

ydb/query/session.py:

def _on_execute_stream_error(self, e: Exception) -> None:
    if isinstance(e, issues.DeadlineExceed):
        self._invalidate()

Nothing else — not Aborted, not BadSession, not Overloaded, not Cancelled, not SessionBusy, not transport errors — invalidates the session. Even if case 1 were fixed and CancelledError reached this method, it still would not invalidate.

End-to-end reproduction

The script below drives the real tx.execute(...)->AsyncResponseContextIterator._next->session._on_execute_stream_error(...) path and compares two exceptions injected at the gRPC-stream layer:

• SessionBusy: the hook is called, but the session stays active.
• DeadlineExceed: the hook is called, and the session is invalidated immediately.

This does not call _on_execute_stream_error directly. The only monkey-patch is on the real async stream iterator used by tx.execute.

# repro_case2_e2e.py  —  ydb==3.26.4
import os

from ydb import issues
from ydb.aio import Driver, QuerySessionPool
from ydb.aio._utilities import AsyncResponseIterator


armed_exc = None
_orig_next = AsyncResponseIterator._next


async def patched_next(self):
    if armed_exc is not None:
        raise armed_exc
    return await _orig_next(self)


AsyncResponseIterator._next = patched_next


async def run_scenario(driver, exc):
    global armed_exc

    victim_sid = None
    hook_events = []

    async with QuerySessionPool(driver, size=1) as pool:
        async def victim(tx):
            nonlocal victim_sid
            global armed_exc
            victim_sid = tx.session.session_id

            orig_hook = tx.session._on_execute_stream_error

            def spy_hook(err):
                before = tx.session.is_active
                result = orig_hook(err)
                after = tx.session.is_active
                hook_events.append((type(err).__name__, before, after))
                return result

            tx.session._on_execute_stream_error = spy_hook

            result = await tx.execute(
                '$rows = ListFromRange(0ul, 100000ul);'
                'SELECT x FROM AS_TABLE(ListMap($rows, ($i) -> (AsStruct($i AS x))));'
            )
            armed_exc = exc
            async for _ in result:
                pass

        try:
            await pool.retry_tx_async(victim)
        except BaseException:
            pass
        armed_exc = None

        async with pool.checkout() as sess:
            reused = sess.session_id == victim_sid
            next_sid = sess.session_id

    print(f'{type(exc).__name__}:')
    print(f'  hook events:          {hook_events}')
    print(f'  next checkout reused: {reused}')
    print(f'  victim sid:           {victim_sid}')
    print(f'  next sid:             {next_sid}')


async def main():
    async with Driver(endpoint=os.environ['YDB_ENDPOINT'], database=os.environ['YDB_DATABASE']) as driver:
        await driver.wait(5, fail_fast=True)
        await run_scenario(driver, issues.SessionBusy('simulated SessionBusy from stream'))
        print()
        await run_scenario(driver, issues.DeadlineExceed('simulated DeadlineExceed from stream'))


if __name__ == '__main__':
    import asyncio
    asyncio.run(main())

Actual output:

SessionBusy:
  hook events:          [('SessionBusy', True, True), ('SessionBusy', True, True)]
  next checkout reused: True
  victim sid:           ydb://session/3?node_id=1&id=...=
  next sid:             ydb://session/3?node_id=1&id=...=

DeadlineExceed:
  hook events:          [('DeadlineExceed', True, False), ('DeadlineExceed', False, False)]
  next checkout reused: False
  victim sid:           ydb://session/3?node_id=1&id=...=
  next sid:             ydb://session/3?node_id=1&id=...=

The contrast is the bug:

• On SessionBusy, the hook runs but leaves the session active (True -> True), so the same session comes back on the next checkout.
• On DeadlineExceed, the same hook invalidates the session (True -> False), and the next checkout gets a different session.

So case 2 is not "the hook never runs" — that is case 1. Case 2 is: even when the hook does run, almost no error types mark the session unusable.

---

Case 3. AsyncResponseContextIterator.__aexit__ drain is not cancel-safe

ydb/aio/query/base.py:

async def __aexit__(self, exc_type, exc_val, exc_tb):
    #  To close stream on YDB it is necessary to scroll through it to the end
    async for _ in self:
        pass

The comment explicitly states that YDB requires draining to the end. The drain is a bare async for on a gRPC aio stream. If the underlying call was cancelled (or anything else went wrong on the wire), __anext__ raises on the very first iteration, the drain exits with zero messages consumed, and __aexit__ re-raises — the server side of the stream stays open. This is what produces the SessionBusy outcome seen in the e2e run above.

---

Case 4. QueryTxContext.__aexit__ skips rollback if drain raises

ydb/aio/query/transaction.py:

async def __aexit__(self, *args, **kwargs):
    await self._ensure_prev_stream_finished()                   # can raise (case 3)
    if self._tx_state._state == QueryTxStateEnum.BEGINED and self._external_error is None:
        ...
        try:
            await self.rollback()                                # never reached under cancel
        except issues.Error:
            logger.warning("Failed to rollback leaked tx: %s", self._tx_state.tx_id)

_ensure_prev_stream_finished() delegates to case 3. The try/except around rollback() only catches issues.Error, so any CancelledError (or AioRpcError) from the drain above propagates out of __aexit__ and rollback() is never called — the e2e run above confirms this with tx.rollback() called: False.

---

Case 5. QuerySessionPool.release puts the session back with no sanity check

ydb/aio/query/pool.py:

async def release(self, session: QuerySession) -> None:
    """Release a session back to Session Pool."""
    self._queue.put_nowait(session)
    logger.debug("Session returned to queue: %s", session.session_id)


class SimpleQuerySessionCheckoutAsync:
    ...
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self._session is not None:
            await self._pool.release(self._session)

There is no "was this session left in a bad state?" check. __aexit__ does not even look at the exception that left the async with — it unconditionally drops the session back into the FIFO queue. The e2e run above confirms this with same session re-queued: True and session.is_active: True, right before the next caller gets SessionBusy.

---

Suggested fixes

These are independent; any subset helps.

A. Do not drop cancelled/errored sessions back into the pool unchecked

The single most effective fix. SimpleQuerySessionCheckoutAsync.__aexit__ should be aware of the exception that left the async with, and on anything unexpected — CancelledError, transport errors, SessionBusy, any non-retriable issues.Error — either:

• call session.delete() / session._invalidate() and let the pool create a new session on the next acquire, or
• perform a synchronous drain + AttachSession health-check before returning the session to the queue.

Current behaviour (unconditional put_nowait) guarantees that one cancelled request poisons a pool slot.

B. Make stream / tx __aexit__ cancel-safe

AsyncResponseContextIterator.__aexit__ and QueryTxContext.__aexit__ should perform their drain/rollback in a cancel-tolerant way. Minimal patch:

async def __aexit__(self, exc_type, exc_val, exc_tb):
    async def _drain():
        async for _ in self:
            pass
    drain = asyncio.ensure_future(_drain())
    try:
        await asyncio.shield(drain)
    except asyncio.CancelledError:
        try:
            await asyncio.wait_for(asyncio.shield(drain), timeout=1.0)
        except (asyncio.CancelledError, asyncio.TimeoutError):
            pass
        raise

A bare async for in __aexit__ of a networked resource is a footgun in any asyncio library.

C. Broaden _on_execute_stream_error and _next error handling

_next should catch BaseException (or at least also invoke _on_error for CancelledError), and _on_execute_stream_error should invalidate the session on CancelledError / transport errors / SessionBusy too. A session whose execute stream was cancelled mid-flight is not safe to reuse.

D. Consider treating SessionBusy as retriable, but only together with session invalidation

There are two competing concerns here:

• On one hand, SessionBusy is a server-side response and clients cannot completely avoid it in all possible situations.
• On the other hand, plain "add SessionBusy to _errors_retriable_*" is not enough and may even hide the real bug if some code path keeps putting broken sessions back into the pool.

In this specific scenario the root cause is a dirty session sitting in the FIFO queue. Naive retry will keep drawing the same dirty session and keep failing until the server-side deadline.

So if SessionBusy is treated as retriable for pool-managed calls, it should be paired with invalidation of the offending session before that session goes back to the pool. With (A) in place, this becomes a natural consequence: any SessionBusy just triggers session invalidation on release.

E. Document the hazard

Until (A)–(D) land, a short note in the QuerySessionPool / retry_tx_async docstrings would save a lot of debugging time:

If the calling task is cancelled while inside this method, the underlying session may be returned to the pool in a bad state. Callers that rely on task.cancel() should either shield calls to this method from cancellation, or be prepared for occasional SessionBusy: Pending previous query completion errors on unrelated subsequent calls for the lifetime of the dirty session.