KMS Retry support - #1999
Conversation
There was a problem hiding this comment.
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
MongoKeyDecryptorAPIs. - 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.
d4e6fbf to
7cad542
Compare
…nd documented the rules in AGENTS.md
There was a problem hiding this comment.
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: 32implies the HTTP body is exactly the JSON string, but this resource currently ends with an extra blank line.getHttpResourceAsByteBufferjoins 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"}
165faec to
cac6f2a
Compare
5ba252a to
a285ac3
Compare
|
Errors in evergreen are related to: DRIVERS-3547 |
2dea3ed to
f985b18
Compare
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
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
left a comment
There was a problem hiding this comment.
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
| } | ||
|
|
||
| // always enabled; backoff and budget management are the driver's responsibility | ||
| configure(() -> mongocrypt_setopt_retry_kms(wrapped, true)); |
There was a problem hiding this comment.
Is this comment useful? For me it brings only confusion.
- "always enabled" - that is pretty obvious, as
mongocrypt_setopt_retry_kmsis 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.
There was a problem hiding this comment.
I can see your confusion - its the responsibility of the driver to sleep. Removed.
| // 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) { |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
I don't think we need to add a commit link the comment is clear enough.
Its all because of the getPointer method returns a null - https://java-native-access.github.io/jna/4.2.1/com/sun/jna/Pointer.html#getPointer-long-
| // 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() { |
There was a problem hiding this comment.
Let's use the documentation comment here, as the comment documents the test.
| // 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() { |
There was a problem hiding this comment.
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.
| protected abstract ClientEncryption getClientEncryption(ClientEncryptionSettings settings); | ||
|
|
||
| @BeforeEach | ||
| public void setUp() { |
There was a problem hiding this comment.
@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.
There was a problem hiding this comment.
This is such a minor nit for test code - I wonder if its even worth doing and reorganising the commits for.
There was a problem hiding this comment.
Done - and added a new rule to .agents/references/testing-guide.md seems the driver is consistently inconsistent here!
It would be good to have these nits automated away via spotless or some other tool in the future.
| 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); | ||
|
|
There was a problem hiding this comment.
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.
| 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 |
There was a problem hiding this comment.
Let's @link to MongoOperationTimeoutException
| 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); |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
Confirmed correct — the retry is in-place at the libmongocrypt level.
The contract is about the native mongocrypt_kms_ctx_t, not the Java wrapper:
feedWithRetrysettingshould_retry(HTTP 429/503) orfail()returningtrue(network error) marks that samekms_ctxfor retry. libmongocrypt keeps it pending rather than done.mongocrypt_ctx_next_kms_ctxthen hands back the samekms_ctx.nextKeyDecryptor()wraps it in a newMongoKeyDecryptorImpl, butwrappedis the same native context — so the retried request is fed to the same context, asintegrating.mdrequires.
This is just libmongocrypt's documented sync loop:
next_kms_ctx → send → feed_with_retry / fail → next_kms_ctx → … → kms_done
The binding docs back this up: feed_with_retry says "retry in-place without calling mongocrypt_kms_ctx_fail", and fail says the ctx "may be reused".
Tests confirm it:
testKmsRetryOnHttpErrorAndNetworkError—reenterNeedKms(decryptor)asserts a non-null decryptor withbytesNeeded() > 0after both a 429 feed and afail(). Ifnext_kms_ctxdidn't re-present retryable contexts, this would fail.testKmsRetryExhaustsBudget— drives repeated in-placefail()retries on the sameMongoKeyDecryptor.
Reactive Crypt.decryptKeys uses the same model (recursing via doOnSuccess), so it behaves identically.
The only non-in-place bit is the per-iteration wrapper allocation — a thin JNA wrapper over the same native pointer, no behavioral effect.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
I'm going to keep them in for now - again we haven't removed CSOT yet and the process of removing it is not set. Eg. deprecation for a major release before another major release or straight removal as its alpha
…r.java Co-authored-by: Valentin Kovalenko <valentin.male.kovalenko@gmail.com>
Co-authored-by: Valentin Kovalenko <valentin.male.kovalenko@gmail.com>
|
Thanks @stIncMale - I've actioned most of the comments. Main open point at the moment is the CSOT implementation as its piped through. I'm the tech lead on DRIVERS-3580 and its yet to go through triage. Also the deprecation / removal process hasn't been determined. So its a cheap add at the moment and we may release this improvement before DRIVERS-3580 is implemented. |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 19 out of 19 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java:64
cstring.toString()only checkspointer == null, butCstringTest(and the Javadoc above) expects it to also treat an explicitPointer.NULLas an empty string. With the current implementation,pointer.getString(0)will be invoked forPointer.NULL, which is likely to return null or error and would failtoStringReturnsEmptyStringForExplicitNullPointer(). Consider normalizing NULL pointers infromNativeand checking for an address of 0 intoString().
public String toString() {
Pointer pointer = getPointer();
return pointer == null ? "" : pointer.getString(0);
}
|
@stIncMale note recent changes in server behaviour are causing some of the CSOT tests to fail. These fail against main as well and are not related to the changes in this PR. See: SERVER-132246 |
Added KMS Retry support.
4 Commits:
JAVA-5391 / JAVA-5772