PYTHON-5672 Refactor connection checkout to reduce layers by blink1073 · Pull Request #2893 · mongodb/mongo-python-driver

blink1073 · 2026-06-25T14:13:01Z

Changes in this PR

Flatten the connection checkout hot path by replacing @asynccontextmanager generator-based context managers with class-based ones, and eliminate _MongoClientErrorHandler by absorbing it into the new _ClientCheckout class. Reduces per-checkout overhead by ~3 generator frames.

Test Plan

just typing passes (mypy + pyright)
just lint-manual passes (all pre-commit hooks)
Existing CMAP and unified format spec tests cover the checkout/checkin contract and SDAM error handling ordering

Checklist

Checklist for Author

Did you update the changelog (if necessary)?
Is there test coverage?
Is any followup work tracked in a JIRA ticket? If so, add link(s). PYTHON-5898 - Reduce server selection and pool lock overhead on the checkout hot path

Checklist for Reviewer

Does the title of the PR reference a JIRA Ticket?
Do you fully understand the implementation? (Would you be comfortable explaining how this code works to someone else?)
Is all relevant documentation (README or docstring) updated?

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-commenter · 2026-06-25T14:51:26Z

Codecov Report

❌ Patch coverage is 88.33333% with 28 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pymongo/asynchronous/mongo_client.py	83.09%	11 Missing and 1 partial ⚠️
pymongo/synchronous/mongo_client.py	83.09%	11 Missing and 1 partial ⚠️
pymongo/asynchronous/pool.py	95.91%	1 Missing and 1 partial ⚠️
pymongo/synchronous/pool.py	95.91%	1 Missing and 1 partial ⚠️

📢 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.

Copilot

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 _PoolCheckout classes (sync + async) and adjust client/server checkout plumbing accordingly.
Introduce _ClientCheckout / _ClientReadCheckout to 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.

…rated files

Copilot

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated 2 comments.

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.

Copilot

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated 2 comments.

aclark4life · 2026-07-23T13:42:57Z

@blink1073 Ready for review ?

blink1073 · 2026-07-23T16:04:20Z

Sure, we can wait to merge until I merge from main and the PR turns green

Merging main brought in _base_backoff_ms (PYTHON-5885) which wasn't present when this branch's __slots__ was written, breaking mypy and causing an AttributeError on every retryable operation.

aclark4life · 2026-07-27T15:23:52Z

+        # connect() already adds cancel_context for new connections; only add
+        # here for reused connections taken from the idle pool.
+        if not is_new_conn:
+            async with self.lock:


🤖 says this is a cancellable await that should move inside the try/except above after 1074 otherwise there is no undo path. With the move the undo path is BaseException and this matches the pre-PR all-or-nothingness of _get_conn.

Thanks, I added a failing test and then fixed it

… reused connection _get_conn() registered a reused connection's cancel_context after its try/except block, so a cancellation during that step skipped the existing cleanup and left active_sockets/requests permanently incremented. Move it inside the try so _get_conn keeps its all-or-nothing contract for reused connections, matching new ones.

aclark4life · 2026-07-27T18:52:43Z

+        )
+
+        async with client.start_session() as session:
+            session.start_transaction()


Missing await

The unawaited coroutine was garbage-collected mid-test, which pytest surfaced as a PytestUnraisableExceptionWarning failure on CI.

Copilot

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated no new comments.

aclark4life

LGTM!

Copilot

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

pymongo/synchronous/pool.py:1088

In _get_conn() the pool increments operation_count before attempting checkout, but the except BaseException: cleanup path does not decrement it. With this refactor (e.g., failures while registering active_contexts for a reused connection), operation_count can leak upward and skew server selection/load comparisons (topology.py uses pool.operation_count). Decrement operation_count alongside requests/active_sockets in the failure cleanup path.

        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
        except BaseException:
            if conn:
                # We checked out a socket but authentication failed.
                conn.close_conn(ConnectionClosedReason.ERROR)
            with self.size_cond:
                self.requests -= 1
                if incremented:
                    self.active_sockets -= 1
                self.size_cond.notify()

pymongo/asynchronous/pool.py:1092

In _get_conn() the pool increments operation_count before attempting checkout, but the except BaseException: cleanup path does not decrement it. With this refactor (e.g., failures while registering active_contexts for a reused connection), operation_count can leak upward and skew server selection/load comparisons (topology.py uses pool.operation_count). Decrement operation_count alongside requests/active_sockets in the failure cleanup path.

        # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
        except BaseException:
            if conn:
                # We checked out a socket but authentication failed.
                await conn.close_conn(ConnectionClosedReason.ERROR)
            async with self.size_cond:
                self.requests -= 1
                if incremented:
                    self.active_sockets -= 1
                self.size_cond.notify()

blink1073 added 2 commits June 25, 2026 09:01

PYTHON-5672 Add __slots__ and lazy _deprioritized_servers to _ClientC…

9b09a4e

…onnectionRetryable Eliminates per-instance __dict__ allocation and defers the _deprioritized_servers list creation until it is actually needed (only on sharded retry paths).

blink1073 added 14 commits June 25, 2026 13:03

PYTHON-5672 Fix MockPool.checkout() bypass in _ClientCheckout

65d0cbe

Constructing _PoolCheckout directly bypassed pool.checkout() overrides used in MockPool, causing mock_down_hosts checks to be skipped and network-error tests to fail.

PYTHON-5672 Fix test_client_checkout_setup_failure on freethreaded Py…

389fc8e

…thon patch.object cannot shadow a method on an instance when __slots__ is defined — the attribute is read-only. Use a subclass override instead.

PYTHON-5672 Merge upstream/master

00bfb71

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.

PYTHON-5672 Remove references to former methods from docstrings

78f1a51

PYTHON-5672 Fix contribute_socket docstring

8554ec0

PYTHON-5672 Move test imports to top of module

ab69293

PYTHON-5672 Merge upstream/master

3247491

PYTHON-5672 Remove implementation detail from _ClientCheckout docstring

a655f01

Merge branch 'main' into PYTHON-5672

3dc0489

blink1073 marked this pull request as ready for review July 22, 2026 11:52

blink1073 requested a review from a team as a code owner July 22, 2026 11:52

blink1073 requested review from aclark4life and Copilot July 22, 2026 11:52

Copilot started reviewing on behalf of blink1073 July 22, 2026 11:52 View session

Copilot AI reviewed Jul 22, 2026

View reviewed changes

blink1073 added 4 commits July 22, 2026 07:55

PYTHON-5672 Fix comments that named __aexit__/__aenter__ in sync-gene…

d45c697

…rated files

Merge remote-tracking branch 'origin/PYTHON-5672' into PYTHON-5672

7fd962a

PYTHON-5672 Add __aenter__/__aexit__ replacements to synchro.py

bd49c2c

PYTHON-5672 Fix copilot review comments

526dc9d

blink1073 requested a review from Copilot July 22, 2026 15:47

Copilot started reviewing on behalf of blink1073 July 22, 2026 15:47 View session

Copilot AI reviewed Jul 22, 2026

View reviewed changes

Comment thread pymongo/synchronous/mongo_client.py Outdated

Comment thread pymongo/asynchronous/mongo_client.py Outdated

blink1073 requested a review from Copilot July 22, 2026 20:04

Copilot started reviewing on behalf of blink1073 July 22, 2026 20:04 View session

Copilot AI reviewed Jul 22, 2026

View reviewed changes

Comment thread pymongo/asynchronous/pool.py

Comment thread pymongo/synchronous/pool.py

blink1073 marked this pull request as draft July 22, 2026 20:38

blink1073 marked this pull request as ready for review July 23, 2026 16:04

aclark4life reviewed Jul 23, 2026

View reviewed changes

Comment thread pymongo/asynchronous/mongo_client.py

blink1073 added 2 commits July 27, 2026 06:29

PYTHON-5672 Add test for session unpin on post-checkout setup failure

28a6e33

PYTHON-5672 Fix test comment wording

5116503

blink1073 requested a review from aclark4life July 27, 2026 13:05

blink1073 added 2 commits July 27, 2026 08:45

Merge remote-tracking branch 'upstream/main' into PYTHON-5672

fbd96d1

PYTHON-5672 Add _base_backoff_ms to _ClientConnectionRetryable __slots__

45d6e6a

Merging main brought in _base_backoff_ms (PYTHON-5885) which wasn't present when this branch's __slots__ was written, breaking mypy and causing an AttributeError on every retryable operation.

aclark4life requested changes Jul 27, 2026

View reviewed changes

blink1073 requested a review from aclark4life July 27, 2026 17:26

aclark4life reviewed Jul 27, 2026

View reviewed changes

blink1073 added 2 commits July 27, 2026 15:05

PYTHON-5672 Add missing await on session.start_transaction() in test

4401bc3

The unawaited coroutine was garbage-collected mid-test, which pytest surfaced as a PytestUnraisableExceptionWarning failure on CI.

Merge remote-tracking branch 'upstream/main' into PYTHON-5672

11c4c5c

blink1073 requested a review from aclark4life July 27, 2026 20:16

aclark4life requested a review from Copilot July 27, 2026 23:12

Copilot started reviewing on behalf of aclark4life July 27, 2026 23:12 View session

Copilot AI reviewed Jul 27, 2026

View reviewed changes

aclark4life requested a review from Copilot July 27, 2026 23:16

Copilot started reviewing on behalf of aclark4life July 27, 2026 23:17 View session

aclark4life approved these changes Jul 27, 2026

View reviewed changes

Copilot AI reviewed Jul 27, 2026

View reviewed changes

blink1073 merged commit e90cc41 into mongodb:main Jul 27, 2026
86 of 89 checks passed

blink1073 deleted the PYTHON-5672 branch July 27, 2026 23:26

Uh oh!

Conversation

blink1073 commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes in this PR

Test Plan

Checklist

Checklist for Author

Checklist for Reviewer

Uh oh!

codecov-commenter commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

aclark4life commented Jul 23, 2026

Uh oh!

blink1073 commented Jul 23, 2026

Uh oh!

Uh oh!

aclark4life Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

blink1073 Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

aclark4life Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

blink1073 Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

aclark4life left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

blink1073 commented Jun 25, 2026 •

edited

Loading

codecov-commenter commented Jun 25, 2026 •

edited

Loading