PYTHON-5672 Refactor connection checkout to reduce layers#2893
PYTHON-5672 Refactor connection checkout to reduce layers#2893blink1073 wants to merge 21 commits into
Conversation
Replace the three @asynccontextmanager layers on the connection checkout hot path with class-based async context managers, and eliminate _MongoClientErrorHandler by absorbing it into _ClientCheckout. - _PoolCheckout replaces Pool.checkout() generator CM - _ClientCheckout replaces _checkout() generator CM and absorbs all of _MongoClientErrorHandler (contribute_socket, handle, SDAM error logic) - _ClientReadCheckout extends _ClientCheckout to apply single-topology read preference adjustment (formerly _conn_from_server()) - active_contexts.add() consolidated into _get_conn(), avoiding a separate lock acquisition on the hot path; deduped so new connections (already tracked by connect()) are not double-added - Connection leak fixed: self._conn assigned before event publishing so checkin runs if a CMAP listener raises in __aenter__ - _ClientCheckout.for_existing_conn() classmethod handles the _run_operation() getMore path that needs SDAM handling around an already-checked-out connection
…onnectionRetryable Eliminates per-instance __dict__ allocation and defers the _deprioritized_servers list creation until it is actually needed (only on sharded retry paths).
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
…nter__ When _PoolCheckout.__aenter__() raises (e.g. connect() fails), Python does not call _ClientCheckout.__aexit__(), so handle() was never invoked and the topology never learned about the failure. This caused test_5_check_out_fails_ connection_error to fail (missing PoolClearedEvent) and broke failover tests. Fix: wrap both the pool checkout call and the post-checkout setup in try/except BaseException blocks inside _ClientCheckout.__aenter__(), calling handle() and (for post-checkout failures) checking the connection back in before re-raising.
Constructing _PoolCheckout directly bypassed pool.checkout() overrides used in MockPool, causing mock_down_hosts checks to be skipped and network-error tests to fail.
…rver.checkout() server.checkout() was a thin wrapper around pool.checkout() that nothing called after _ClientCheckout was changed to call server.pool.checkout() directly. Add a test that mocks publish_connection_checked_out to raise and verifies the connection is returned to the pool rather than leaked.
…connection Covers the except BaseException block in _ClientCheckout.__aenter__() that fires when post-checkout work (session pinning, ConfigurationError check) raises — verifying the connection is returned to the pool and not leaked.
…thon patch.object cannot shadow a method on an instance when __slots__ is defined — the attribute is read-only. Use a subclass override instead.
Resolve import conflict in server.py: upstream removed datetime import (run_operation timing moved into run_cursor_command), our branch removed AbstractAsyncContextManager (server.checkout() was deleted). Drop both now-unused imports.
- Remove _checkout_started_time from _PoolCheckout.__slots__: the value was written then immediately aliased to a local variable and never read from the slot again, wasting an allocation on every checkout. - Guard pool_checkout.__aexit__ with try/finally in both _ClientCheckout __aexit__ and the __aenter__ setup-failure path: if handle() raises, the connection is now guaranteed to be checked in instead of leaked. - Pass exc.__traceback__ instead of None when calling pool_checkout.__aexit__ from the __aenter__ setup-failure except block.
- Reset conn.pinned_txn/pinned_cursor before calling pool_checkout.__aexit__ in the setup-failure except block so _PoolCheckout.__aexit__ calls checkin() instead of silently skipping it when session._pin() ran before the error. - Replace _ClientReadCheckout._effective_read_pref slot with _read_preference and a local variable in __aenter__ to avoid mutating instance state. - Fix Pool.checkout() docstring to say "with-statement" so synchro.py copies the correct text to the sync mirror.
There was a problem hiding this comment.
Pull request overview
Refactors PyMongo’s connection checkout hot path to reduce overhead by replacing generator-based @contextmanager / @asynccontextmanager usage with class-based context managers, and by folding _MongoClientErrorHandler responsibilities into a new _ClientCheckout context manager for both sync and async clients.
Changes:
- Replace pool checkout generator context managers with
_PoolCheckoutclasses (sync + async) and adjust client/server checkout plumbing accordingly. - Introduce
_ClientCheckout/_ClientReadCheckoutto unify pool checkout, SDAM error handling, and session pinning (sync + async). - Add regression tests to ensure connections are not leaked when CMAP event publishing fails or when post-checkout setup raises.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| test/test_read_preferences.py | Updates read-preference test client override to match new checkout return types. |
| test/test_pooling.py | Adds regression test ensuring checkout doesn’t leak connections if CMAP listener publish fails. |
| test/test_client.py | Adds regression test ensuring client checkout setup failures return connections to the pool. |
| test/pymongo_mocks.py | Updates MockPool checkout override to match new class-based pool checkout behavior. |
| test/asynchronous/test_read_preferences.py | Async equivalent of read-preference override update. |
| test/asynchronous/test_pooling.py | Async equivalent of CMAP publish-failure no-leak regression test. |
| test/asynchronous/test_client.py | Async equivalent of client checkout setup-failure no-leak regression test. |
| test/asynchronous/pymongo_mocks.py | Async equivalent of MockPool checkout override update. |
| pymongo/synchronous/server.py | Removes now-unneeded server-level checkout wrapper in favor of direct pool checkout usage. |
| pymongo/synchronous/pool.py | Implements class-based _PoolCheckout and adjusts pool checkout internals for reuse/new-connection paths. |
| pymongo/synchronous/mongo_client.py | Introduces _ClientCheckout / _ClientReadCheckout and re-routes operation execution through them. |
| pymongo/asynchronous/server.py | Removes now-unneeded async server-level checkout wrapper. |
| pymongo/asynchronous/pool.py | Async implementation of class-based _PoolCheckout. |
| pymongo/asynchronous/mongo_client.py | Async implementation of _ClientCheckout / _ClientReadCheckout and updated operation execution flow. |
When session._pin() runs before ConfigurationError is raised during __aenter__, clear the session's pinned state (pinned_address, conn_mgr) in the except block so future operations don't route to a stale server or double-checkin via conn_mgr.
|
@blink1073 Ready for review ? |
|
Sure, we can wait to merge until I merge from main and the PR turns green |
| conn.pinned_cursor = False | ||
| # Clear stale session pin state so future ops don't reuse it. | ||
| if session and session.in_transaction: | ||
| await session._unpin() |
There was a problem hiding this comment.
How about a test to confirm that _unpin happens when some post-checkout step fails e.g. https://github.com/mongodb/mongo-python-driver/pull/2893/changes#diff-8dfbbdba94166c58103fd44781f9133a77767f13c9b816030a631f29d99aee5bR2743 ?
PYTHON-5672
Changes in this PR
Flatten the connection checkout hot path by replacing
@asynccontextmanagergenerator-based context managers with class-based ones, and eliminate_MongoClientErrorHandlerby absorbing it into the new_ClientCheckoutclass. Reduces per-checkout overhead by ~3 generator frames.Test Plan
just typingpasses (mypy + pyright)just lint-manualpasses (all pre-commit hooks)Checklist
Checklist for Author
Checklist for Reviewer