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

blink1073 · 2025-08-08T21:56:27Z

This has a couple commits from #2467

blink1073 · 2025-08-08T22:07:06Z

Okay this is ready for a look. We might consider switching to the base Protocol since the sizes of the responses are constrained and small, and we typically only have 2-3 reads per socket.

blink1073 · 2025-08-08T22:07:32Z

I still need to make a PR for the fix to the kms mock server's 404 response.

blink1073 · 2025-08-09T01:14:58Z

CSOT failure is unrelated: PYTHON-5492

NoahStapp · 2025-08-11T15:30:18Z

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 sizehint = -1, which signals that the buffer size can be arbitrary, this check will always succeed, potentially returning an empty buffer, which is an error. We need to check that sizehint is a positive number as well.

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.

NoahStapp · 2025-08-11T15:45:38Z

pymongo/network_layer.py

+        """
+        self.transport = transport  # type: ignore[assignment]
+
+    async def read(self, bytes_needed: int) -> bytes:


To make sure I understand the intended flow here, is this example correct?

We call kms_request and enter the while kms_context.bytes_needed > 0: loop.

The first chunk of data, say 16 bytes worth is written into an existing buffer that still has space.

PyMongoKMSProtocol.read() is called and immediately returns those 10 bytes.

kms_context.bytes_needed updates to need 84 more bytes for a total of 100.

We call PyMongoKMSProtocol.read() again and wait on the _pending_listeners Future we create.

The second chunk of data, the remaining 84 bytes, requires a new buffer since the active buffer is full.

The Future is resolved with those bytes, which we return and feed into kms_context to complete the operation.

We call kms_request and enter the while kms_context.bytes_needed > 0: loop.

The first chunk of data, say 16 bytes worth is written into an existing buffer that still has space.

PyMongoKMSProtocol.read() is called and immediately returns those 10 bytes, pushing the start_index up by 10

kms_context.bytes_needed updates to need 84 more bytes

We call PyMongoKMSProtocol.read() again and wait on the _pending_listeners Future we created.

The second chunk of data, the remaining 84 bytes, may require a new buffer

If any bytes are available, we read in up to the newly requested 84 bytes from the active buffer(s), advancing start_index and exhausting buffers as appropriate

Otherwise, we wait on the future to be resolved, which will contain up to the requested bytes.

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

blink1073 requested a review from NoahStapp August 8, 2025 22:07

NoahStapp requested changes Aug 11, 2025

View reviewed changes

blink1073 added 2 commits August 11, 2025 10:52

address review

4546f23

use upstream d-e-t

39b4526

blink1073 marked this pull request as ready for review August 11, 2025 15:53

blink1073 requested a review from NoahStapp August 11, 2025 15:53

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


		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

Uh oh!

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!

blink1073 commented Aug 8, 2025

Uh oh!

blink1073 commented Aug 8, 2025

Uh oh!

blink1073 commented Aug 8, 2025

Uh oh!

blink1073 commented Aug 9, 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!

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!

Uh oh!

blink1073 commented Aug 6, 2025 •

edited

Loading

NoahStapp left a comment •

edited

Loading