fix(memory): enforce closed state in MongoDBSession - #4134
Conversation
MongoDBSession.close() closed an owned client but never recorded that the session was closed, so every session method kept issuing commands against a closed client instead of failing fast. get_items() would still return items and add_items() would still return None, leaving the caller with no signal that the write went nowhere. Applies the pattern already merged for the other backends (openai#4035 for Redis and Dapr, openai#4109 for AsyncSQLiteSession): a _closed flag plus a _check_not_closed() helper, checked in _ensure_indexes() (the chokepoint all four protocol methods funnel through) and explicitly in add_items and get_items, which can early-return before reaching it. ping() is guarded too. close() is terminal from the first call and idempotent, closing an owned client exactly once and leaving an injected client untouched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7a8445c1e2
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| already_closed = self._closed | ||
| self._closed = True | ||
| if self._owns_client and not already_closed: | ||
| await self._client.close() |
There was a problem hiding this comment.
With a caller-supplied AsyncMongoClient, this now marks the session closed before checking ownership, so await session.close() makes later get_items()/add_items() fail even though the client is still open. That breaks the documented MongoDB lifecycle contract in docs/sessions/index.md:451 that session.close() is a no-op when the application owns the client, and it also diverges from the Redis/Dapr injected-client behavior; only owned from_uri(...) sessions should become terminal here.
AGENTS.md reference: AGENTS.md:L84-L87
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in 9b932a7 — thanks, this was a real contract break.
docs/sessions/index.md:451 does say session.close() is a no-op for an injected client, and RedisSession.close() / DaprSession.close() both return before touching _closed when not self._owns_client. My version set _closed first, so an injected-client session became terminal even though its client stayed open.
close() now returns early for injected clients, matching Redis/Dapr:
if not self._owns_client:
return
self._closed = TrueI had also written a test asserting the wrong contract (that an injected-client session becomes terminal). It has been replaced by test_close_injected_client_is_a_noop, which asserts the documented behaviour instead — the client stays open and get_items / add_items / ping all keep working after close(). That test fails against the previous commit.
| already_closed = self._closed | ||
| self._closed = True | ||
| if self._owns_client and not already_closed: | ||
| await self._client.close() |
There was a problem hiding this comment.
For an owned from_uri(...) session, if await self._client.close() raises or is cancelled, _closed has already been set and every later close() call takes the already_closed no-op path, so the owned MongoDB client can be left open with no way for the session to retry cleanup. This failure/cancellation path should track whether client release actually completed, or otherwise avoid suppressing later close attempts after unsuccessful cleanup.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in 9b932a7.
You're right that _closed = True before the release completed made the failure unrecoverable: the next close() took the no-op path and the owned client stayed open forever. RedisSession and DaprSession already solve this with a separate _client_released flag, so I've adopted it:
if not self._owns_client:
return
self._closed = True
if not self._client_released:
await self._client.close()
self._client_released = TrueThe session is still terminal from the first attempt (operations raise), but an unfinished release is retried. test_failed_owned_client_close_can_be_retried covers it with a client whose first close() raises ConnectionError:
| step | before 9b932a7 | after |
|---|---|---|
1st close() |
raises ConnectionError, _closed=True |
raises ConnectionError, _closed=True |
get_items() |
RuntimeError |
RuntimeError |
2nd close() |
no-op — client left open | retries release, client closed |
3rd close() |
no-op | no-op (release already done) |
That test fails against the previous commit too.
…etry Addresses two review findings on the previous commit: - With a caller-supplied client, close() marked the session closed before checking ownership, making later operations raise even though the client was still open. That contradicts the documented contract in docs/sessions/index.md that close() is a no-op when the application owns the client, and diverged from Redis/Dapr. close() now returns early for injected clients, so the session never becomes terminal. - For an owned client, _closed was set before the release completed, so a failed or cancelled client.close() left the client open with no way to retry. Release is now tracked separately via _client_released, matching RedisSession and DaprSession: the session is terminal from the first attempt, but a later close() retries unfinished cleanup. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9b932a78d9
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if not self._client_released: | ||
| await self._client.close() | ||
| self._client_released = True |
There was a problem hiding this comment.
When two coroutines call close() on the same owned MongoDBSession while AsyncMongoClient.close() is still awaiting I/O, both can pass this _client_released guard before either call sets it. The new idempotence state therefore still allows duplicate cleanup, so a second close can run against an already-closing client or surface a cleanup error even though another close succeeds; the added gather() test misses this because the fake close() never suspends. Guard this section with an async lock or shared close task.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in 13ebaf4 — and you're right that my gather() test couldn't have caught it.
The fake client's close() returned without ever suspending, so the second coroutine never got to run while the first was inside the release. Making it await asyncio.sleep(0) mid-release reproduces the double release immediately:
client.close() invocations under concurrent close(): 2 (before 13ebaf4)
client.close() invocations under concurrent close(): 1 (after)
test_close_is_idempotent_and_closes_owned_client_once now suspends mid-release for exactly this reason, and it fails against 9b932a7 with assert 2 == 1.
One wrinkle worth flagging on the remedy: unlike RedisSession/DaprSession, this class cannot hold an asyncio.Lock created in __init__. Its _init_state design notes explicitly call out that a session may be used from more than one event loop, and an asyncio.Lock binds to the loop that first acquires it. So the lock is created lazily and rebound if close() runs on a different loop:
def _close_guard(self) -> AbstractAsyncContextManager[Any]:
with self._init_guard:
loop = asyncio.get_running_loop()
if self._close_lock is None or self._close_lock_loop is not loop:
self._close_lock = asyncio.Lock()
self._close_lock_loop = loop
return self._close_lockRebinding is safe because a differing loop means no concurrent close is in flight on the previous one. test_close_from_a_second_event_loop_does_not_error covers that path (it closes the session from a thread running its own loop), so the serialization fix doesn't quietly reintroduce the cross-loop restriction the class was written to avoid.
_closed is still set before awaiting the lock, so a caller queued behind a slow release still can't issue commands in the meantime.
46 tests pass; ruff, mypy and pyright are clean.
…ed once Concurrent close() calls could both pass the _client_released guard while AsyncMongoClient.close() was still awaiting I/O, releasing the client twice. The previous idempotence test could not catch this because the fake client's close() never suspended. close() now holds a lock around the release. The lock is created lazily and rebound if the session is closed from a different event loop than the one that created it: an asyncio.Lock binds to a single loop, and this session is documented as usable across loops (the same reason _init_state is guarded by a threading.Lock). _closed is still set before awaiting the lock, so a caller blocked behind a slow release cannot issue commands in the meantime. The idempotence test now suspends mid-release, which reproduces the double release against the previous commit, and a new test closes a session from a second event loop. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 13ebaf47b3
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if self._close_lock is None or self._close_lock_loop is not loop: | ||
| self._close_lock = asyncio.Lock() | ||
| self._close_lock_loop = loop |
There was a problem hiding this comment.
When two event loops close the same owned session concurrently, the second loop takes this branch and replaces the in-flight lock with a fresh unlocked one, so both callers can enter the if not self._client_released block before either awaits AsyncMongoClient.close() to completion. The same-loop gather() test now passes, but a shutdown path that calls close() from different threads can still release the owned client twice despite the new docstring promising concurrent calls release it exactly once; use a loop-independent guard/shared close task instead of rebinding the lock while a close may be in flight.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in e02c68e — you're right, and the flaw was in the reasoning I put in the docstring, not just the code.
I justified rebinding with "a differing loop means no concurrent close is in flight on the previous one." That's exactly the assumption two threads violate. Reproduced with two threads running their own loops, released behind a threading.Barrier, with a release that stays in-flight for 50ms:
client.close() invocations across two loops: 2 (before e02c68e)
client.close() invocations across two loops: 1 (after)
Rather than rebinding, the release is now claimed under the threading.Lock this class already uses — which serializes callers regardless of loop, and is the same reason _init_state is guarded by a threading lock rather than an asyncio.Lock:
def _claim_client_release(self) -> bool:
with self._init_guard:
if self._client_released or self._releasing_client:
return False
self._releasing_client = True
return Trueclose() returns early if it loses the claim, and a try/finally calls _finish_client_release(released) so failure or cancellation frees the claim without marking the client released — the retry behaviour from 9b932a7 is preserved. No asyncio.Lock is involved anywhere now, so there's nothing left to rebind. _closed is still set before the claim, so a caller that loses it still can't issue commands.
test_concurrent_close_across_event_loops_releases_client_once covers this and fails against 13ebaf4 with assert 2 == 1. The same-loop gather() test (which fails against 9b932a7) is kept, since the two paths reach the guard differently.
47 tests pass; ruff, mypy and pyright clean.
… event loop The previous fix rebound the asyncio.Lock when close() ran on a different loop, which defeated the purpose: two loops (or threads) closing the same session would each install a fresh unlocked lock and both enter the release, so the owned client was released twice despite the docstring promising once. Replaced with a claim/finish pair guarded by the existing threading.Lock, which serializes callers regardless of which loop they run on - the same reason _init_state uses a threading.Lock rather than an asyncio.Lock. The claim only guards flag flips, so no async coordination is needed. Release failures and cancellation free the claim without marking the client released, preserving the retry behaviour. test_concurrent_close_across_event_loops_releases_client_once closes from two threads with their own loops and fails against the previous commit with assert 2 == 1. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e02c68e4a4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| return | ||
| # Mark terminal before releasing, so a concurrent caller that returns | ||
| # early below still cannot issue commands against a closing client. | ||
| self._closed = True |
There was a problem hiding this comment.
Because this change makes an owned MongoDBSession terminal on the first close(), the live MongoDB docs are now incomplete: docs/sessions/index.md:451 still only says from_uri(...) closes the client and injected-client close() is a no-op, while the Redis/Dapr sections tell users that post-close operations raise and repeated/concurrent closes are safe. Please update the English session docs alongside this SDK behavior change so MongoDB users do not learn a different lifecycle from the published docs.
AGENTS.md reference: AGENTS.md:L86-L86
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in 42c77db. The Redis and Dapr bullets both already state this and the MongoDB one didn't, so the docs were inconsistent as well as incomplete:
docs/sessions/index.md:367 (Redis) "the session is terminal and subsequent session operations raise RuntimeError; repeated or concurrent close() calls are safe"
docs/sessions/index.md:417 (Dapr) "makes an owned-client session terminal; subsequent session operations raise RuntimeError, while repeated or concurrent close() calls are safe"
docs/sessions/index.md:451 (Mongo) -- only ownership and the injected-client no-op
The MongoDB bullet now covers the same three points in the same order: owned-client sessions are terminal after close() with operations raising RuntimeError, repeated/concurrent closes are safe and release the client exactly once, and with an injected client close() is a no-op leaving both lifecycle and session usability with the caller.
Only the English docs are edited — the translations are generated, so I've left them to make translate_docs.
| return | ||
| # Mark terminal before releasing, so a concurrent caller that returns | ||
| # early below still cannot issue commands against a closing client. | ||
| self._closed = True |
There was a problem hiding this comment.
Because this change makes an owned MongoDBSession terminal on the first close(), the live MongoDB docs are now incomplete: docs/sessions/index.md:451 still only says from_uri(...) closes the client and injected-client close() is a no-op, while the Redis/Dapr sections tell users that post-close operations raise and repeated/concurrent closes are safe. Please update the English session docs alongside this SDK behavior change so MongoDB users do not learn a different lifecycle from the published docs.
Useful? React with 👍 / 👎.
| if not self._claim_client_release(): | ||
| return |
There was a problem hiding this comment.
When an owned-client close() is already awaiting AsyncMongoClient.close() and that in-flight release later fails or is cancelled, this fast path lets every other concurrent await close() return successfully without waiting for cleanup or observing the failure. Fresh evidence after the earlier fixes is that non-claiming callers now take this immediate return, so a shutdown path can finish with _client_released still false and the Mongo client still open until some later close attempt happens to run.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in 42c77db — this one was a genuine weakening versus the sibling backends, and I'd missed why.
In RedisSession/DaprSession the losing caller blocks on the asyncio.Lock and then re-tests _client_released, so it retries a failed release. My claim-based version returned immediately instead, which meant a caller could get a clean return from a release that raised. Reproduced with two concurrent closes where the release fails:
before 42c77db after
close() #1 (claims) ConnectionError ConnectionError
close() #2 (loses) returned None <-- bug ConnectionError
client released no no
later close() retries retries
Losing callers now await the in-flight attempt:
claimed, pending = self._claim_client_release()
if not claimed:
if pending is not None:
await asyncio.wrap_future(pending)
returnThe attempt is a concurrent.futures.Future, not an asyncio one, so it can be awaited from any loop via asyncio.wrap_future — keeping the cross-loop property from e02c68e rather than trading it back for this fix. _finish_client_release publishes the exception to every waiter and leaves _client_released false, so a later close() still retries.
test_concurrent_close_observes_a_failed_release pins it: both gather() results must be the ConnectionError, the client must still be open, and a third close() must succeed. It fails against e02c68e.
48 tests pass; ruff, mypy and pyright clean; full suite unchanged against the merge base (59 vs 59, the delta being the known flaky guardrail/tracing tests).
…ycle Two follow-ups from review: - A caller that lost the release claim returned immediately, so if the in-flight release then failed or was cancelled, that caller reported success while the client stayed open. Losing callers now await the in-flight attempt and observe its failure, matching Redis/Dapr where losers wait on the lock. The attempt is a concurrent.futures.Future awaited via asyncio.wrap_future, so it still works across event loops - an asyncio future would bind the session to one loop. - docs/sessions/index.md documented only that from_uri owns the client and that injected-client close() is a no-op. The Redis and Dapr sections already state the terminal behaviour, so the MongoDB bullet now says the same: post-close operations raise, repeated and concurrent closes are safe and release the client exactly once, and an injected client leaves both lifecycle and session usability with the caller. test_concurrent_close_observes_a_failed_release fails against the previous commit, where one of two concurrent callers returned successfully from a release that raised. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 42c77dbf77
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if pending is not None: | ||
| # Await the in-flight release rather than returning early, so a | ||
| # failed cleanup surfaces here too instead of looking successful. | ||
| await asyncio.wrap_future(pending) |
There was a problem hiding this comment.
When one owned-client close() is in flight, a second caller that is cancelled while awaiting asyncio.wrap_future(pending) cancels the shared concurrent.futures.Future; the releasing caller then hits InvalidStateError in _finish_client_release() when it tries to publish the successful/failed cleanup result. This can happen in shutdown paths that wrap concurrent closes in timeouts or cancel sibling tasks, and it makes the primary close report a spurious error even after the Mongo client was released; wrap the pending future in a cancellation shield or otherwise prevent waiter cancellation from cancelling the shared release attempt.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in 65466c6.
asyncio.wrap_future forwards cancellation to the wrapped concurrent.futures.Future, so a cancelled waiter destroyed the attempt the releasing caller still needed. Reproduced with a slow release, a second close() task cancelled mid-await:
Before (42c77dbf):
RELEASER BROKE: InvalidStateError CANCELLED: <Future state=cancelled>
After (65466c62):
releaser completed OK; client closed = True
Waiters now go through a helper that shields the wrapper, so cancellation stops at the waiter and never reaches the shared attempt:
@staticmethod
async def _await_client_release(pending: Future[None]) -> None:
wrapper = asyncio.wrap_future(pending)
try:
await asyncio.shield(wrapper)
except asyncio.CancelledError:
if not wrapper.done():
wrapper.add_done_callback(lambda f: f.cancelled() or f.exception())
raiseThe done callback is there for the second half of the problem: an abandoned wrapper left holding an exception (the failed-release path from the previous review) would otherwise emit "exception was never retrieved" at GC time, so its outcome is consumed on the way out.
Added test_cancelled_waiter_does_not_break_the_in_flight_release, which asserts the releaser still completes, the client is released, and the session remains terminal afterwards. It fails against 42c77dbf with exactly the InvalidStateError: CANCELLED above and passes with the fix. 49 passed in the file; ruff, ruff format, mypy and pyright clean on both changed files.
A caller waiting on another close()'s in-flight release awaited the shared concurrent.futures.Future directly through asyncio.wrap_future, which forwards cancellation to the wrapped future. Cancelling the waiter therefore cancelled the shared attempt, and the releasing caller crashed with InvalidStateError when it tried to publish its outcome. Shield the wrapper so a cancelled waiter cannot abort a release it does not own, and discard the abandoned wrapper's outcome so a failed release does not log an unretrieved exception. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 65466c62ac
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if not wrapper.done(): | ||
| wrapper.add_done_callback(lambda f: f.cancelled() or f.exception()) |
There was a problem hiding this comment.
After the shielding fix, a waiter that is cancelled at the same time the shared close attempt fails can enter this CancelledError handler with wrapper.done() already true; this branch then skips both the callback and wrapper.exception(), so asyncio logs the Mongo close exception as “Future exception was never retrieved” despite the docstring promising to discard abandoned outcomes. This shows up in shutdown or timeout paths that cancel a secondary owned-client close() while the release is failing; consume the exception when the wrapper is already done too.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Investigated this one carefully and the mechanism is not quite as described, but it pointed at a real leak, so the fix is in 73f88d5a.
The wrapper.done() branch is not itself a leak: asyncio.shield's own _inner_done_callback consumes the inner exception when the outer future is already cancelled.
def _inner_done_callback(inner):
if outer.cancelled():
if not inner.cancelled():
# Mark inner's result as retrieved.
inner.exception()
returnWhat decides the outcome is the ordering, because _outer_done_callback removes that inner callback once the outer is cancelled. Measured both interleavings with the callback stripped out:
cancel delivered, then release fails -> logged=['Future exception was never retrieved']
release fails, then cancel delivered -> logged=[]
So the leak lives in the not done() case that my guard was already covering — but the guard was pointless rather than protective, and reasoning about which of two callbacks wins a race is not something worth keeping in the code. The callback is now added unconditionally: asyncio invokes it immediately for an already-completed future, and re-reading a retrieved outcome is harmless.
test_cancelled_waiter_consumes_a_failed_release is parametrized over both orderings and asserts the loop exception handler stays silent. The cancel_first case fails when the callback is removed and passes with it; fail_first passes either way, which is exactly the asymmetry above.
Marking this useful — the conclusion about wrapper.done() didn't hold up, but it sent me to the shield internals, and the simplification is a genuine improvement.
| self._closed = True | ||
| claimed, pending = self._claim_client_release() |
There was a problem hiding this comment.
When an owned-client close() races with a session operation that already passed _check_not_closed() (for example get_items() after _ensure_indexes() or add_items() while building the payload), this marks the session terminal and proceeds to close the AsyncMongoClient without waiting for that in-flight operation. That operation can then resume and issue its MongoDB command against a client that has already been closed, so shutdown paths can still get PyMongo closed-client failures or partial writes instead of the new RuntimeError contract; serialize close with active session methods or add an equivalent lifecycle gate around the actual Mongo commands.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in 73f88d5. This one was a real gap in the port: all three already-fixed backends serialize close() against operations, and I had only carried over the flag.
redis_session.py every method: async with self._lock: self._check_not_closed() close(): async with self._lock
dapr_session.py same shape (lock at :106, close at :441)
async_sqlite_session.py _locked_connection() holds self._lock; close(): async with self._lock
mongodb_session.py -- flag only, no serialization
Reproduced by suspending the collection command after the guard had already passed, then closing underneath it:
Before (65466c62): command ran against a CLOSED client: [True]
After (73f88d5a): command ran against a CLOSED client: [False]
I did not copy the lock, though — asyncio.Lock is exactly what this class cannot use, for the reason documented on _init_state and raised in your earlier review: it binds to the loop that first acquires it, and this session is documented as usable from more than one loop. Holding one across every operation would also newly serialize concurrent reads and writes that overlap freely today.
Instead operations register on a counter guarded by the existing threading.Lock, and close() drains it after setting _closed (so nothing new can enter) and before claiming the release:
@contextlib.contextmanager
def _operation(self) -> Iterator[None]:
with self._init_guard:
self._check_not_closed()
self._active_operations += 1
try:
yield
finally:
... # signal _idle when the count reaches zeroThe drain awaits a concurrent.futures.Future through the same shielded helper, so it works from any loop and a cancelled close() cannot abort a drain other callers are waiting on. add_items([]) still checks the flag directly, since it returns before doing any work.
test_close_waits_for_an_in_flight_operation asserts the command sees an open client, the client is still released, and the session is terminal afterwards. It fails against 65466c62 with At index 0 diff: True != False.
On the sibling finding about consuming completed close failures — that one I could not reproduce as described, but investigating it did surface a real bug next door, so the fix is in.
asyncio.shield's _inner_done_callback already consumes the inner exception when the outer future is cancelled, so the wrapper.done() path is not itself a leak. What actually matters is the ordering, because _outer_done_callback removes that inner callback on cancel:
cancel delivered, then release fails -> logged=['Future exception was never retrieved']
release fails, then cancel delivered -> logged=[]
So the leak is in the not done() case, and my if not wrapper.done() guard was the wrong shape for a different reason than stated — it was dead weight rather than protection. The callback is now added unconditionally (asyncio invokes it immediately for an already-completed future, and re-reading a retrieved outcome is harmless), and test_cancelled_waiter_consumes_a_failed_release is parametrized over both interleavings. The cancel_first case fails without the callback and passes with it.
52 passed in the file, 290 across tests/extensions/memory; ruff, ruff format, mypy and pyright clean on both changed files.
…ient An operation that had already passed the closed check could resume after close() released the client and issue its command against a closed AsyncMongoClient, producing PyMongo errors or partial writes instead of the documented RuntimeError. The other backends avoid this by holding a lock across both the operation and close(); an asyncio.Lock is unusable here because this session is documented as usable from more than one event loop. Track active operations under the threading.Lock this class already uses and have close() drain them before releasing, so concurrent reads and writes still overlap and only close() waits. Also consume an abandoned waiter's outcome unconditionally: asyncio.shield only marks the wrapped exception retrieved when the failure lands before the cancellation is delivered. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
PranavMishra28
left a comment
PranavMishra28
left a comment
There was a problem hiding this comment.
The threading design reads well. concurrent.futures.Future plus threading.Lock rather than asyncio primitives is the right call for a class documented as cross-loop, and _operation() closes the check-then-act window properly.
One case worth handling: cancelling close() propagates a cancel into tasks that were never cancelled. At mongodb_session.py:546-553 the except BaseException catches asyncio.CancelledError, and _finish_client_release then reaches attempt.set_exception (:514) with it, publishing a CancelledError onto the shared future. Waiters in _await_client_release (:479-497) sit on await asyncio.shield(wrapper) over asyncio.wrap_future, and asyncio's _convert_future_exc maps a concurrent.futures.CancelledError back to asyncio.CancelledError, so line 495 raises it in a task nobody cancelled and the handler at :496-497 re-raises.
I reproduced the asyncio half standalone, mirroring _await_client_release exactly:
waiter task done : True
waiter task CANCELLED : True
waiter exception : t.exception() itself raised CancelledError
So the waiter's task ends up marked cancelled rather than failed. Its own finally still runs, but the real error is invisible to the caller and anything sequenced after that await is skipped, which for a session close is usually the rest of a teardown path.
Would translating it be preferable to propagating, something like attempt.set_exception(RuntimeError("MongoDB client close was cancelled; retry close()")), or leaving _release_attempt = None so waiters retry the release instead of inheriting the cancel? A test where close() is cancelled mid-client.close() while a second close() waits, asserting the waiter's task is not cancelled, would pin whichever you pick.
Cancelling the `close()` that owns the release published the resulting `CancelledError` onto the shared attempt future. `asyncio` maps a `concurrent.futures.CancelledError` back to its own on the way out of `wrap_future`, so every waiter's task ended up marked *cancelled* rather than *failed*: the real cause was invisible to the caller and anything the teardown path sequenced after the await was silently skipped. Publish a `RuntimeError` instead. Only the caller that was actually cancelled sees `CancelledError`, which `close()` re-raises itself. The client is still left unreleased, so a later `close()` retries. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Thanks — you're right, and the reproduction holds end-to-end in the session, not just in the standalone asyncio mirror. Fixed in Probe: releaser blocks inside I went with your first suggestion — translate rather than propagate: elif isinstance(error, asyncio.CancelledError):
attempt.set_exception(
RuntimeError("MongoDB client release was cancelled; call close() again to retry")
)I preferred it over leaving Cancellation semantics after the fix: only the caller that was actually cancelled sees Regression: Full suite vs the branch base ( |
|
Thanks for trying to fix this. The issue is going to be resolved by #4176 |
3 participants
Summary
MongoDBSession.close()closes the underlying client when the session owns it (created viafrom_uri), but it never records that the session is closed. Every session method then keeps issuing commands against a closed client instead of failing fast.This is the same gap that was closed for the other backends:
RedisSession/DaprSessionAsyncSQLiteSessionMongoDBSessionwas the remaining backend with an unguardedclose().Repro
Against
main, with an owned client (from_uri) afterawait session.close():The same script against the already-fixed backend:
get_itemsreturning items andadd_itemsreturningNoneafter the client is closed is the worst shape of this: the caller has no signal that the write went nowhere.Fix
Mirrors the merged pattern exactly:
self._closed = Falsein__init__, set toTrueinclose()._check_not_closed()helper raisingRuntimeError("MongoDBSession is closed")._operation()context manager, so all four methods andping()are covered at a single point.add_itemschecks the flag directly on its empty-list fast path, which returns before doing any work.ping()is guarded too, matching how fix(memory): enforce closed state in Redis and Dapr sessions #4035 guarded Redis'sping._active_operationscounter guarded by that samethreading.Lock, andclose()drains it after marking the session closed and before releasing the client, so an operation that already passed the guard cannot issue its command against a closed client. The other backends get this by holding anasyncio.Lockacross both the operation andclose(), which is not available here; the counter also keeps concurrent reads and writes overlapping as they do today, with onlyclose()waiting.close()follows the Redis/Dapr shape: it returns early for an injected client (a no-op, per the documented contract that lifecycle stays with the caller), and for an owned client the session is terminal from the first attempt while_client_releasedtracks whether the release actually completed, so a failed or cancelledclient.close()is retried by a laterclose(). The release is claimed under thethreading.Lockthis class already uses for_init_state, not anasyncio.Lock, so concurrent closes are serialized even across event loops and threads — this session is documented as usable from more than one loop, and anasyncio.Lockbinds to whichever loop first acquires it.After the fix:
The MongoDB bullet in
docs/sessions/index.mdis also updated: it documented only client ownership, while the Redis and Dapr bullets already state the terminal-after-close behaviour. Translations are generated, so they are left tomake translate_docs.Tests
Fifteen tests added to
tests/extensions/memory/test_mongodb_session.py, following the convention from #4109:pingraise afterclose()add_items([])does not slip past the guard via the empty-list fast pathget_items(limit=0)does not slip past it via the non-positive-limit fast pathclose()before any command still makes the session terminalclose()are no-ops and the owned client is closed exactly once (the fake client suspends mid-release, without which a same-loop race is unobservable)close()does not abort the release it was waiting on (the shared attempt is shielded, so the releasing caller can still publish its outcome)close()drains an operation that already passed the guard instead of releasing the client underneath itclose()from a second event loop does not raiseclose()rather than being suppressedclose()stays a no-op and the session keeps workingAll six guard tests fail on
main(DID NOT RAISE <class 'RuntimeError'>) and pass with the fix; the negative control passes on both.pytest tests/extensions/memory/test_mongodb_session.py→ 52 passed (290 acrosstests/extensions/memory).ruff check,ruff format,mypyandpyrightare clean on both changed files, and the full suite shows no change against the merge base.