PYTHON-5215 Add an asyncio.Protocol implementation for KMS #2460

blink1073 · 2025-08-06T19:30:40Z

See benchmark gist.

Benchmark Results:

Before: 4.93s, 5.26s
After: 4.93s, 5.05s

Depends on mongodb-labs/drivers-evergreen-tools#679

blink1073 · 2025-08-06T20:14:51Z

I'm debugging two failures:

test.asynchronous.test_connection_monitoring.AsyncTestCMAP.test_connection_monitoring_pool_clear_interrupting_pending_connections_clear_with_interruptInUseConnections___true_closes_pending_connections
test.asynchronous.test_connection_monitoring.AsyncTestCMAP.test_connection_monitoring_pool_create_min_size_error_error_during_minPoolSize_population_clears_pool

…PYTHON-5215

blink1073 · 2025-08-07T00:25:36Z

I realized the benchmark test wasn't actually triggering the protocol -> I'm tweaking things locally

…PYTHON-5215

NoahStapp · 2025-08-11T16:02:34Z

pymongo/network_layer.py


-async def async_sendall(conn: PyMongoProtocol, buf: bytes) -> None:
+        bytes_needed = self._pending_reads.popleft()


What happens if we need more bytes than we have? We've already popped the waiter and set it's result to data, which can only read up to self._bytes_ready bytes. Are we relying on the kms_context.bytes_needed loop to call the protocol read() method again and create a new waiter?

Right, we give the partial result back to the kms context, and let it ask for more.

Maybe we could get better performance by doing more of the looping inside the Protocol, but KMS requests won't be a significant part of runtime anyway so not worth spending more time on it. Can you add a comment to this effect somewhere saying that we rely on the looping behavior for this to function correctly?

It's not really a question of perf, but the fact that the kms_request is blind until it knows the Content-Length, and we don't know what state it is in.

I added a comment.

NoahStapp · 2025-08-14T18:15:15Z

test/asynchronous/test_collection.py

@@ -335,6 +335,8 @@ async def test_create_index(self):
        await db.test.create_index(["hello", ("world", DESCENDING)])
        await db.test.create_index({"hello": 1}.items())  # type:ignore[arg-type]

+    # TODO: PYTHON-5491 - remove version max


This change should be in a separate PR.

Yeah it was, I just updated this branch.

NoahStapp · 2025-08-14T18:18:15Z

pymongo/network_layer.py

+        # Reuse the active buffer if it has space.
+        if len(self._buffers):
+            buffer = self._buffers[-1]
+            if len(buffer.buffer) - buffer.end_index > sizehint:


If we're setting sizehint to be at least 16384 always, is this check worth doing in the first place? I'd expect us to rarely reuse the active buffer since we'll usually have a buffer of size 16384 and a sizehint of 16384.

blink1073 · 2025-08-14T18:59:05Z

If we're setting sizehint to be at least 16384 always, is this check worth doing in the first place? I'd expect us to rarely reuse the active buffer since we'll usually have a buffer of size 16384 and a sizehint of 16384.

The actual sizehint in practice was on the order of the bytes being read from the buffer (typically less than 1000). Using the buffered protocol at all here is a bit of a mismatch imho.

NoahStapp · 2025-08-14T19:42:21Z

If we're setting sizehint to be at least 16384 always, is this check worth doing in the first place? I'd expect us to rarely reuse the active buffer since we'll usually have a buffer of size 16384 and a sizehint of 16384.

The actual sizehint in practice was on the order of the bytes being read from the buffer (typically less than 1000). Using the buffered protocol at all here is a bit of a mismatch imho.

How long would refactoring to not use buffered take? No reason to use the lower-level API if we don't need to.

blink1073 · 2025-08-14T21:23:31Z

How long would refactoring to not use buffered take? No reason to use the lower-level API if we don't need to.

It's actually dead simple, I did it along the way when I was debugging a race condition.

blink1073 · 2025-08-14T21:24:41Z

It's actually dead simple, I did it along the way when I was debugging a race condition.

I'll push a commit in the morning for comparison, we can always revert.

…r into PYTHON-5215

blink1073 · 2025-08-15T14:22:09Z

I'm happy with the simplification. The tests are passing locally, this is ready for another look.

NoahStapp

Can you schedule a full Evergreen run? We should ensure there's no regressions introduced here by accident.

Are the benchmark results for KMS significantly different between the two Protocol implementations?

…PYTHON-5215

blink1073 · 2025-08-15T15:39:27Z

Full patch build: https://spruce.mongodb.com/version/689f5483e112170007b0ce9f/tasks?sorts=STATUS%3AASC%3BBASE_STATUS%3ADESC

I updated the timings in the PR description, there no significant change.

blink1073 · 2025-08-15T16:04:24Z

Okay there is one legit bug in test.test_connection_logging.TestConnectionLoggingConnectionLogging.test_Connection_checkout_fails_due_to_error_establishing_connection. I'll defer looking at that until next week to focus on greener build tasks.

…PYTHON-5215

blink1073 · 2025-08-18T17:58:28Z

Here's a new patch build, failures are existing flakiness issues or tracked in PYTHON-5502.

https://spruce.mongodb.com/version/68a3434b33ee5e00072767c9/tasks?sorts=STATUS%3AASC%3BBASE_STATUS%3ADESC

NoahStapp · 2025-08-18T18:51:35Z

pymongo/synchronous/pool.py

@@ -124,7 +124,89 @@ def _set_non_inheritable_non_atomic(fd: int) -> None:  # noqa: ARG001
 _IS_SYNC = True


-class Connection:
+class BaseConnection:


We don't utilize any of the BaseConnection abstractions for the sync API, correct? The change is purely for compatibility with the async API?

No it is used by both sync and async KMS, it is what gets returned by the _connect_kms helper function.

I meant that we don't have a separate KMS networking interface for sync. Just wanted to confirm my understanding.

…PYTHON-5215

blink1073 · 2025-08-19T00:00:03Z

I had to fix merge conflicts

PYTHON-5215 Add an asyncio.Protocol implementation for KMS

a490cfd

blink1073 requested a review from NoahStapp August 6, 2025 19:30

blink1073 requested a review from a team as a code owner August 6, 2025 19:30

blink1073 added 3 commits August 6, 2025 14:31

cleanup

404c1fc

restore comment

e4a588b

fix close

2af62a3

blink1073 added 2 commits August 6, 2025 15:17

Merge branch 'master' of github.com:mongodb/mongo-python-driver into …

8494954

…PYTHON-5215

wip

0a43477

blink1073 marked this pull request as draft August 7, 2025 00:24

blink1073 removed the request for review from NoahStapp August 7, 2025 00:24

blink1073 added 18 commits August 6, 2025 19:29

wip

13537cb

wip

18c51cb

wip

24aa733

wip

b33f78e

fixup

4b1bdd6

always allow partial reads

fa0dd8d

fixup

432380e

cleanup

484aa9f

fix sync kms

e50685c

undo change to justfile

2622a7a

remove unused code

73b4309

undo lock file changes

0cbdd58

use det branch

6fe6ba3

fix branch name

6ed92bb

fix buffer handling and close handling

971139c

Merge branch 'master' of github.com:mongodb/mongo-python-driver into …

4f174f9

…PYTHON-5215

fix close conn behavior

db4332d

skip another test

c2f6ae8

NoahStapp requested changes Aug 11, 2025

View reviewed changes

blink1073 added 2 commits August 11, 2025 12:37

fix waiting logic

da04fc8

address review

28afc38

blink1073 requested a review from NoahStapp August 11, 2025 21:09

NoahStapp requested changes Aug 14, 2025

View reviewed changes

Merge branch 'master' into PYTHON-5215

33add41

blink1073 requested a review from NoahStapp August 14, 2025 19:00

blink1073 added 3 commits August 14, 2025 20:04

use the base Protocol

8b357cd

fixups

0708278

Merge branch 'PYTHON-5215' of github.com:blink1073/mongo-python-drive…

14a6199

…r into PYTHON-5215

NoahStapp reviewed Aug 15, 2025

View reviewed changes

Merge branch 'master' of github.com:mongodb/mongo-python-driver into …

513e66a

…PYTHON-5215

blink1073 requested a review from NoahStapp August 15, 2025 15:39

blink1073 added 2 commits August 18, 2025 07:21

PYTHON-5500 Mark test_dns_failures_logging as flaky

b5e8c0e

Merge branch 'master' of github.com:mongodb/mongo-python-driver into …

2a3d4c3

…PYTHON-5215

NoahStapp reviewed Aug 18, 2025

View reviewed changes

blink1073 requested a review from NoahStapp August 18, 2025 20:44

NoahStapp previously approved these changes Aug 18, 2025

View reviewed changes

Merge branch 'master' of github.com:mongodb/mongo-python-driver into …

3edba6f

…PYTHON-5215

blink1073 dismissed NoahStapp’s stale review via 3edba6f August 18, 2025 23:59

blink1073 requested a review from NoahStapp August 18, 2025 23:59


		async def async_sendall(conn: PyMongoProtocol, buf: bytes) -> None:
		bytes_needed = self._pending_reads.popleft()

PYTHON-5215 Add an asyncio.Protocol implementation for KMS #2460

Are you sure you want to change the base?

PYTHON-5215 Add an asyncio.Protocol implementation for KMS #2460

Conversation

blink1073 commented Aug 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

blink1073 commented Aug 6, 2025

Uh oh!

blink1073 commented Aug 7, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

blink1073 commented Aug 14, 2025

Uh oh!

NoahStapp commented Aug 14, 2025

Uh oh!

blink1073 commented Aug 14, 2025

Uh oh!

blink1073 commented Aug 14, 2025

Uh oh!

blink1073 commented Aug 15, 2025

Uh oh!

NoahStapp left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

blink1073 commented Aug 15, 2025

Uh oh!

blink1073 commented Aug 15, 2025

Uh oh!

blink1073 commented Aug 18, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

blink1073 commented Aug 19, 2025

Uh oh!

Uh oh!

blink1073 commented Aug 6, 2025 •

edited

Loading

NoahStapp left a comment •

edited

Loading