KMS Retry support by rozza · Pull Request #1999 · mongodb/mongo-java-driver

rozza · 2026-06-11T11:39:59Z

Added KMS Retry support.

4 Commits:

Update CAPI javadoc and added rules to: ‎mongodb-crypt/AGENTS.md‎
Implement KMS Retry API in mongo-crypt
Implemented KMS Retries in sync / reactive streams
CAPI cstring NPE hardening - found by copilot

Copilot

Pull request overview

This PR adds end-to-end support for libmongocrypt’s KMS retry behavior across the sync and reactive-streams drivers, and updates the vendored TLS channel to correctly surface close_notify when it arrives in the same network segment as application data (addressing the KMS HTTP response parsing edge case referenced by JAVA-5391 / JAVA-5411).

Changes:

Enable and surface libmongocrypt KMS retry primitives (retryable HTTP errors, network failures, backoff sleep) via new JNA bindings and MongoKeyDecryptor APIs.
Update sync/reactive KMS I/O loops to use retry-aware feeding (feedWithRetry), budget exhaustion (fail), and pre-request sleep (sleepMicroseconds), including CSOT-aware timeout behavior.
Add prose/integration tests and Evergreen wiring for KMS retry scenarios, plus a unit test suite for the TLS close_notify + data coalescing behavior.

Reviewed changes

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

Show a summary per file

File	Description
mongodb-crypt/src/test/resources/kms-reply-429.txt	Adds a canned HTTP 429 KMS response used to validate retryable HTTP error handling.
mongodb-crypt/src/test/java/com/mongodb/crypt/capi/MongoCryptTest.java	Adds unit tests validating retry-on-HTTP-error, retry-on-network-error, and retry budget exhaustion.
mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoKeyDecryptorImpl.java	Implements new retry/backoff APIs by calling newly bound libmongocrypt KMS functions.
mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoKeyDecryptor.java	Extends the decryptor interface with sleep/backoff and retry signaling methods.
mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoCryptImpl.java	Enables libmongocrypt KMS retry behavior during initialization.
mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java	Adds new JNA native method declarations for KMS retry and expands/updates binding documentation.
mongodb-crypt/AGENTS.md	Documents contribution rules and mapping guidelines for `CAPI.java` JNA bindings/javadoc.
driver-sync/src/test/functional/com/mongodb/client/ClientSideEncryptionKmsRetryProseTest.java	Adds sync-driver entrypoint for the KMS retry prose tests.
driver-sync/src/test/functional/com/mongodb/client/AbstractClientSideEncryptionKmsRetryProseTest.java	Implements KMS retry prose tests (network retry, HTTP retry, budget exhaustion, CSOT timing).
driver-sync/src/main/com/mongodb/client/internal/Crypt.java	Integrates retry-aware KMS decryption into the sync encryption state machine, including backoff and timeout mapping.
driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/ClientSideEncryptionKmsRetryProseTest.java	Adds reactive-streams entrypoint for the KMS retry prose tests.
driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/KeyManagementService.java	Adds reactive retry/backoff handling and retry-aware KMS response feeding.
driver-core/src/test/unit/com/mongodb/internal/connection/tlschannel/TlsChannelCloseNotifyTest.java	Adds regression tests for TLS channel read semantics when `close_notify` and data coalesce.
driver-core/src/main/com/mongodb/internal/connection/tlschannel/impl/TlsChannelImpl.java	Updates vendored TLS-channel read/unwrap behavior to return bytes produced alongside `close_notify` and eagerly unwrap buffered records.
.evergreen/run-kms-retry-tests.sh	Adds a CI script to run the new sync + reactive KMS retry prose tests.
.evergreen/.evg.yml	Wires a new Evergreen function/task/buildvariant to execute the KMS retry test script.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…nd documented the rules in AGENTS.md

Copilot

Pull request overview

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

Comments suppressed due to low confidence (1)

mongodb-crypt/src/test/resources/kms-reply-429.txt:7

Content-Length: 32 implies the HTTP body is exactly the JSON string, but this resource currently ends with an extra blank line. getHttpResourceAsByteBuffer joins lines with \r\n, so the trailing empty line appends an extra CRLF after the body, making the response longer than the declared Content-Length and potentially breaking/parsing the retry logic in a non-representative way.

Remove the trailing blank line so the file ends immediately after the JSON body (matching how kms-reply.txt is formatted).

Content-Length: 32

{"__type":"ThrottlingException"}

Copilot

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 1 comment.

Copilot

Pull request overview

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

rozza · 2026-06-17T13:37:09Z

Errors in evergreen are related to: DRIVERS-3547

Expose libmongocrypt's KMS retry primitives so transient KMS failures can be retried with library-supplied backoff delays. - Add CAPI bindings: mongocrypt_setopt_retry_kms, mongocrypt_kms_ctx_usleep, mongocrypt_kms_ctx_feed_with_retry and mongocrypt_kms_ctx_fail - Expose them via sleepMicroseconds(), feedWithRetry() and fail() on MongoKeyDecryptor, and enable retry_kms unconditionally in MongoCryptImpl - HTTP 429 responses are signalled by feedWithRetry and network errors by fail(); in both cases the attempt ends early and libmongocrypt re-presents the KMS context for retry with the next backoff delay - Add mongodb-crypt unit tests for the new bindings and a 429 KMS reply fixture JAVA-5391

Wire libmongocrypt's KMS retry behaviour through the sync and reactive drivers. Transient KMS failures are retried with backoff delays supplied by libmongocrypt; retry is enabled unconditionally. - The sync Crypt sleeps the backoff before each attempt, surfaces MongoOperationTimeoutException only once the CSOT deadline has actually expired (a fixed-timeout connect or handshake stall with budget remaining stays retryable) and treats unexpected EOF as a retryable network error - The reactive KeyManagementService applies the backoff via Mono.delay, checks the CSOT deadline before consuming retry budget, and releases buffers and closes streams on all completion paths; the shared backoff budget check lives in MongoCryptHelper - Add the spec Section 24 KMS retry prose tests (sync and reactive) plus a CSOT-during-retry test; each test resets the failpoint server and control requests send Connection: close so the single-threaded server cannot stall KMS connections or leak armed failpoints between tests and suites - Add a sync unit test covering the CSOT classification of KMS network failures, and an Evergreen task running the KMS retry suites JAVA-5391

Copilot

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 1 comment.

rozza · 2026-06-17T16:18:59Z

    /**
-     * Get the error message associated with a status, or an empty string.
+     * Get the error message associated with a status or NULL.
     *
     * @param status The status object.
-     * @param len an optional length of the returned string. May be NULL.
-     * @return An error message or an empty string.
+     * @param len    An optional length of the returned string (excluding the
+     *               trailing NULL byte). May be NULL.
+     * @return A NULL terminated error message or NULL.
     */
    public static native cstring
    mongocrypt_status_message(mongocrypt_status_t status, Pointer len);


Fixed - in an separate commit.

Although this has been here since pre merge to this repo. Its a simple hardening fix to ensure NPE's dont propgate to the user. I can move this out of this review if required.

Native functions declared to return a cstring (e.g. mongocrypt_status_message) may return a NULL pointer. By default JNA marshals that to a null cstring reference, so call sites that immediately call toString() would throw a NullPointerException. Override cstring.fromNative to return a cstring with a null pointer rather than a null reference, and have toString normalize a null pointer to an empty string. This makes the binding's documented NULL contract safe for existing callers without changing each call site.

stIncMale

This is a partial review. I will come back to reviewing this PR once we are out of the woods with the backpressure project.

The previous PR was closed and replaced with this one. Let's not do the same to this PR (also, let's not force push to the PR branch), as that generally complicates the review process.

Review status of individual files as of 38a268fd7abd774a246a68023c5467ade21ce86c (✓ means the file was reviewed, ✗ means it wasn't):

✓ .evergreen/.evg.yml
✓.evergreen/run-kms-retry-tests.sh
✓ driver-core/src/main/com/mongodb/internal/capi/MongoCryptHelper.java
✗ driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/KeyManagementService.java
✓ driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/ClientSideEncryptionKmsRetryProseTest.java
✗ driver-reactive-streams/src/test/unit/com/mongodb/reactivestreams/client/internal/crypt/KeyManagementServiceKmsRetryTest.java
✗ driver-sync/src/main/com/mongodb/client/internal/Crypt.java
✓ driver-sync/src/test/functional/com/mongodb/client/AbstractClientSideEncryptionKmsRetryProseTest.java
✓ driver-sync/src/test/functional/com/mongodb/client/ClientSideEncryptionKmsRetryProseTest.java
✗ driver-sync/src/test/unit/com/mongodb/client/internal/CryptKmsRetryTest.java
✓ mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java - I did not review the unrelated API documentation changes, and do not plan to review them.
✓ mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoCryptImpl.java
✓ mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoKeyDecryptor.java
✓ mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoKeyDecryptorImpl.java
✓ mongodb-crypt/src/test/java/com/mongodb/crypt/capi/MongoCryptTest.java
✓ mongodb-crypt/src/test/java/com/mongodb/internal/crypt/capi/CstringTest.java
✓ mongodb-crypt/src/test/resources/kms-reply-429.txt
✓ mongodb-crypt/AGENTS.md

stIncMale · 2026-07-17T19:27:21Z

        }

+        // always enabled; backoff and budget management are the driver's responsibility
+        configure(() -> mongocrypt_setopt_retry_kms(wrapped, true));


Is this comment useful? For me it brings only confusion.

"always enabled" - that is pretty obvious, as mongocrypt_setopt_retry_kms is called unconditionally. What useful information does this part add?

"backoff and budget management are the driver's responsibility" - as far as I understand, whether to retry and how long to sleep before retrying is managed by lybmongocrypt, not the driver, while the comment seem to imply the opposite.

stIncMale · 2026-07-20T19:03:45Z

+        // By default JNA returns a null cstring reference when a native function returns a NULL
+        // pointer (e.g. mongocrypt_status_message for a status with no message). Override fromNative
+        // to return a cstring with a null pointer instead, so callers never receive a null reference
+        // and toString below can normalize it to an empty string.
+        @Override
+        public Object fromNative(final Object nativeValue, final FromNativeContext context) {


I did not find the JNA behavior mentioned in the comment anywhere in the JNA documentation. Let's add the java-native-access/jna@3fc6bc7 commit link to the comment as the supporting evidence.

stIncMale · 2026-07-20T20:21:32Z

+    // A native function declared to return a cstring (e.g. mongocrypt_status_message) may return a
+    // NULL pointer. By default JNA would marshal that to a null cstring reference, causing callers
+    // that immediately call toString() to throw a NullPointerException. The fromNative override must
+    // return a non-null cstring instead.
+    @Test
+    void fromNativeReturnsEmptyStringBackedInstanceForNullPointer() {


Let's use the documentation comment here, as the comment documents the test.

stIncMale · 2026-07-20T20:29:20Z

+    // A native function declared to return a cstring (e.g. mongocrypt_status_message) may return a
+    // NULL pointer. By default JNA would marshal that to a null cstring reference, causing callers
+    // that immediately call toString() to throw a NullPointerException. The fromNative override must
+    // return a non-null cstring instead.
+    @Test
+    void fromNativeReturnsEmptyStringBackedInstanceForNullPointer() {


This test method name is misleading. fromNative(null, null) does not return an "empty string–backed instance", as can be observed by result.getPointer().getString(0) throwing NullPointerException. That is, the returned cstring is backed by Pointer.NULL, but its toString method returns the empty String.

stIncMale · 2026-07-21T15:27:35Z

+    protected abstract ClientEncryption getClientEncryption(ClientEncryptionSettings settings);
+
+    @BeforeEach
+    public void setUp() {


@BeforeEach, @AfterEach, @ParameterizedTest, @Test methods do not have to be public. Let's make them all package-access.

I implemented the proposed change in stIncMale@189a6d6, tested here.

stIncMale · 2026-07-23T12:32:08Z

+        keyDecryptor = reenterNeedKms(decryptor);
+        assertTrue(keyDecryptor.bytesNeeded() > 0);
+        assertTrue(keyDecryptor.getMessage().remaining() > 0);
+
+        // A network error is retryable while budget remains
+        assertTrue(keyDecryptor.fail());
+        keyDecryptor = reenterNeedKms(decryptor);
+


Let's move the following two assertTrue invocations

assertTrue(keyDecryptor.bytesNeeded() > 0); assertTrue(keyDecryptor.getMessage().remaining() > 0);

into reenterNeedKms? This way, we will neither miss them on the second quoted invocation of reenterNeedKms, nor will have to duplicate them there.

stIncMale · 2026-07-23T12:54:00Z

+    public static final String KMS_TIMEOUT_ERROR_MESSAGE = "KMS key decryption exceeded the timeout limit.";
+
+    /**
+     * Throws a {@code MongoOperationTimeoutException} if the operation timeout has expired or the


Let's @link to MongoOperationTimeoutException

stIncMale · 2026-07-23T15:03:54Z

+        assertTrue(keyDecryptor.feedWithRetry(getHttpResourceAsByteBuffer("kms-reply-429.txt")));
+        keyDecryptor.sleepMicroseconds();
+
+        // The context stays in NEED_KMS and the key decryptor is re-presented for the retry
+        keyDecryptor = reenterNeedKms(decryptor);
+        assertTrue(keyDecryptor.bytesNeeded() > 0);
+        assertTrue(keyDecryptor.getMessage().remaining() > 0);


[informational comment]

The PR description says that it implements JAVA-5772: Support retrying a KMS in-place (and indeed, the PR uses the MongoKeyDecryptor.feedWithRetry (mongocrypt_kms_ctx_feed_with_retry(mongocrypt_kms_ctx_t, ...)) method from that ticket, yet reenterNeedKms calls MongoCryptContext.nextKeyDecryptor (mongocrypt_ctx_next_kms_ctx(mongocrypt_ctx_t)) after feeding a KMS error response. This surprised me, as in an in-place retry "If a KMS context indicates retry, retry the KMS request and feed the new response to the same KMS context".

com.mongodb.client.internal.Crypt.decryptKeys also does not do an in-place retry (at the moment, I am unsure about com.mongodb.reactivestreams.client.internal.crypt.Crypt.decryptKeys) - it calls MongoCryptContext.nextKeyDecryptor each time it retries.

Interestingly, testKmsRetryExhaustsBudget does reuse the MongoKeyDecryptor (mongocrypt_kms_ctx_t), thus testing an in-place retry of a KMS network error.

stIncMale · 2026-07-23T15:07:54Z

+        operationTimeout.run(TimeUnit.MICROSECONDS,
+                // infinite timeout: no CSOT budget to enforce; libmongocrypt's retry count is the only limit
+                () -> { },
+                remainingMicros -> {
+                    if (remainingMicros < backoffMicros) {
+                        throw TimeoutContext.createMongoTimeoutException(KMS_TIMEOUT_ERROR_MESSAGE);
+                    }
+                },
+                () -> {
+                    throw TimeoutContext.createMongoTimeoutException(KMS_TIMEOUT_ERROR_MESSAGE);
+                });


Suggested change

operationTimeout.run(TimeUnit.MICROSECONDS,

// infinite timeout: no CSOT budget to enforce; libmongocrypt's retry count is the only limit

() -> { },

remainingMicros -> {

if (remainingMicros < backoffMicros) {

throw TimeoutContext.createMongoTimeoutException(KMS_TIMEOUT_ERROR_MESSAGE);

}

},

() -> {

throw TimeoutContext.createMongoTimeoutException(KMS_TIMEOUT_ERROR_MESSAGE);

});

operationTimeout.shortenBy(backoffMicros, TimeUnit.MICROSECONDS).onExpired(() -> {

throw TimeoutContext.createMongoTimeoutException(KMS_TIMEOUT_ERROR_MESSAGE);

});

stIncMale · 2026-07-23T15:26:40Z

The changes in this file are all CSOT-related, and @dariakp confirmed that we should not do any more CSOT-related work (except for removing CSOT, or making sure the code still compiles and works). Thus, I think we should remove all the CSOT-related changes from the PR, including CSOT-related tests in MongoCryptTest, KeyManagementServiceKmsRetryTest.

rozza requested a review from Copilot June 11, 2026 11:40

Copilot started reviewing on behalf of rozza June 11, 2026 11:40 View session

Copilot AI reviewed Jun 11, 2026

View reviewed changes

Comment thread mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java

Comment thread mongodb-crypt/AGENTS.md Outdated

rozza force-pushed the JAVA-5391 branch 3 times, most recently from d4e6fbf to 7cad542 Compare June 11, 2026 14:06

Updated CAPI.java javadocs to their 1.18.1 mongocrypt.h equivalents a…

df2a759

…nd documented the rules in AGENTS.md

rozza force-pushed the JAVA-5391 branch from 7cad542 to e5e4b77 Compare June 11, 2026 14:08

rozza requested a review from Copilot June 11, 2026 14:11

Copilot started reviewing on behalf of rozza June 11, 2026 14:11 View session

rozza changed the title ~~Java 5391~~ KMS Retry support Jun 11, 2026

Copilot AI reviewed Jun 11, 2026

View reviewed changes

rozza force-pushed the JAVA-5391 branch 3 times, most recently from 165faec to cac6f2a Compare June 15, 2026 11:03

rozza requested a review from Copilot June 15, 2026 11:24

Copilot started reviewing on behalf of rozza June 15, 2026 11:25 View session

Copilot AI reviewed Jun 15, 2026

View reviewed changes

Comment thread ...streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/KeyManagementService.java Outdated

rozza force-pushed the JAVA-5391 branch 2 times, most recently from 5ba252a to a285ac3 Compare June 16, 2026 09:29

rozza requested a review from Copilot June 16, 2026 13:41

Copilot started reviewing on behalf of rozza June 16, 2026 13:42 View session

Copilot AI reviewed Jun 16, 2026

View reviewed changes

Comment thread ...streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/KeyManagementService.java

Comment thread mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java

rozza force-pushed the JAVA-5391 branch 2 times, most recently from 2dea3ed to f985b18 Compare June 17, 2026 14:11

rozza force-pushed the JAVA-5391 branch from f985b18 to 428571a Compare June 17, 2026 14:31

rozza force-pushed the JAVA-5391 branch from 428571a to 68115e8 Compare June 17, 2026 14:45

rozza requested a review from Copilot June 17, 2026 14:51

Copilot started reviewing on behalf of rozza June 17, 2026 14:51 View session

Copilot AI reviewed Jun 17, 2026

View reviewed changes

rozza marked this pull request as ready for review June 18, 2026 13:36

rozza requested a review from a team as a code owner June 18, 2026 13:36

rozza requested a review from vbabanin June 18, 2026 13:36

rozza force-pushed the JAVA-5391 branch from c0f67ef to 226b0f7 Compare July 6, 2026 12:14

rozza added 3 commits July 7, 2026 10:59

Merge branch 'main' into JAVA-5391

bf0140f

Merge branch 'main' into JAVA-5391

3f6a2ce

Merge branch 'main' into JAVA-5391

4a5c864

stIncMale requested review from stIncMale and removed request for vbabanin July 16, 2026 00:27

Merge branch 'main' into JAVA-5391

38a268f

stIncMale requested changes Jul 23, 2026

View reviewed changes

Uh oh!

Conversation

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

Uh oh!

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!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

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!

rozza commented Jun 17, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

rozza Jun 17, 2026 • 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!

stIncMale left a comment

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

rozza commented Jun 11, 2026 •

edited

Loading

rozza Jun 17, 2026 •

edited

Loading