Do not leak query pool capacity on cancelled session create - #872
Open
vgvoleg wants to merge 2 commits into
Open
Do not leak query pool capacity on cancelled session create#872vgvoleg wants to merge 2 commits into
vgvoleg wants to merge 2 commits into
Conversation
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
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #872 +/- ##
==========================================
+ Coverage 82.33% 82.37% +0.03%
==========================================
Files 99 99
Lines 12749 12752 +3
Branches 1242 1242
==========================================
+ Hits 10497 10504 +7
+ Misses 1797 1794 -3
+ Partials 455 454 -1
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
Fixes an async QuerySessionPool.acquire() cancellation edge case where pool capacity could be leaked permanently, eventually causing all subsequent acquires to block forever. The PR also hardens both sync and async session attach paths so interrupted/cancelled attach attempts properly tear down the stream and invalidate the session to avoid server-side orphans.
Changes:
- Add
except BaseExceptionrollback in asyncQuerySessionPool.acquire()to restore_current_sizeon cancellation/interruption during session creation. - Broaden
_attach()guards (sync and async) fromExceptiontoBaseExceptionto ensure interrupted attach always closes/invalidates the session and cancels the stream. - Add regression tests covering acquire-cancellation capacity leaks and attach interruption/cancellation teardown, plus a user-facing changelog entry.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| ydb/aio/query/pool.py | Roll back _current_size on BaseException during session creation to prevent leaked pool capacity. |
| ydb/aio/query/pool_test.py | Add async regression tests for cancelled create/acquire and cancelled attach behavior. |
| ydb/aio/query/session.py | Ensure cancelled/interrupting attach invalidates session and tears down attach stream. |
| ydb/query/session.py | Ensure interrupted attach invalidates session and tears down attach stream. |
| ydb/query/pool_test.py | Add sync regression tests for interrupted attach teardown behavior. |
| CHANGELOG.md | Document the user-visible fix for async pool cancellation capacity leaks and attach teardown. |
Suppressed comments (2)
ydb/aio/query/pool_test.py:202
- This
await entered.wait()has no timeout, so a regression could cause the test run to hang indefinitely. Preferasyncio.wait_forwith a bounded timeout to fail fast.
await entered.wait()
ydb/aio/query/pool_test.py:221
- This
await entered.wait()has no timeout, so if the attach call never signals the event the test can hang indefinitely. Wrap it inasyncio.wait_forto bound the wait.
await entered.wait()
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
140
to
+144
| # 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 | ||
| except BaseException: |
| pool._create_new_session = self._hanging_create(entered) | ||
|
|
||
| task = asyncio.create_task(pool.acquire()) | ||
| await entered.wait() |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #870.
asyncio.CancelledErrorderives fromBaseException, so theexcept Exceptionrollback in asyncQuerySessionPool.acquire()never ran when a task was cancelled inside_create_new_session()._current_sizestayed incremented and the slot was lost for good; once the counter reachedsize, everyacquire()blocked on_queue.get()forever.Added an
except BaseExceptionbranch that restores the counter. Also widened the_attach()guard (sync and async) fromExceptiontoBaseException, so a cancelled/interrupted attach closes the session and cancels the stream instead of orphaning it server-side.