Skip to content

MONGOCRYPT-763 add in-place retry API #1011

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 9 commits into
base: master
Choose a base branch
from
Open

MONGOCRYPT-763 add in-place retry API #1011

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
ce988fd
implement
mdb-ad May 17, 2025
f615153
format
mdb-ad May 17, 2025
4cb5326
Apply suggestions from code review
mdb-ad Jun 18, 2025
ab662af
Combine fail and reset
mdb-ad Jun 18, 2025
db1b417
describe combined retry API
mdb-ad Jun 18, 2025
ac957c9
Merge branch 'master' into retry-inplace
mdb-ad Jun 18, 2025
63b36a2
document in-place retry
mdb-ad Jul 3, 2025
7e44b5f
remove unused kms_ctx_reset
mdb-ad Jul 3, 2025
d2951a4
Merge branch 'master' into retry-inplace
mdb-ad Jul 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 30 additions & 17 deletions src/mongocrypt-kms-ctx.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1120,6 +1120,35 @@ static bool _ctx_done_kmip_decrypt(mongocrypt_kms_ctx_t *kms_ctx) {
return ret;
}

static bool _is_retryable_req(_kms_request_type_t req_type) {
// Check if request type is retryable. Some requests are non-idempotent and cannot be safely retried.
_kms_request_type_t retryable_types[] = {MONGOCRYPT_KMS_AZURE_OAUTH,
MONGOCRYPT_KMS_GCP_OAUTH,
MONGOCRYPT_KMS_AWS_ENCRYPT,
MONGOCRYPT_KMS_AWS_DECRYPT,
MONGOCRYPT_KMS_AZURE_WRAPKEY,
MONGOCRYPT_KMS_AZURE_UNWRAPKEY,
MONGOCRYPT_KMS_GCP_ENCRYPT,
MONGOCRYPT_KMS_GCP_DECRYPT};
for (size_t i = 0; i < sizeof(retryable_types) / sizeof(retryable_types[0]); i++) {
if (retryable_types[i] == req_type) {
return true;
}
}
return false;
}

bool mongocrypt_kms_ctx_should_retry_http(mongocrypt_kms_ctx_t *kms) {
return kms && kms->should_retry;
}

void mongocrypt_kms_ctx_reset(mongocrypt_kms_ctx_t *kms) {
if (kms->parser) {
kms_response_parser_reset(kms->parser);
}
kms->should_retry = false;
}

bool mongocrypt_kms_ctx_fail(mongocrypt_kms_ctx_t *kms) {
if (!kms) {
return false;
Expand All @@ -1138,23 +1167,7 @@ bool mongocrypt_kms_ctx_fail(mongocrypt_kms_ctx_t *kms) {
return false;
}

// Check if request type is retryable. Some requests are non-idempotent and cannot be safely retried.
_kms_request_type_t retryable_types[] = {MONGOCRYPT_KMS_AZURE_OAUTH,
MONGOCRYPT_KMS_GCP_OAUTH,
MONGOCRYPT_KMS_AWS_ENCRYPT,
MONGOCRYPT_KMS_AWS_DECRYPT,
MONGOCRYPT_KMS_AZURE_WRAPKEY,
MONGOCRYPT_KMS_AZURE_UNWRAPKEY,
MONGOCRYPT_KMS_GCP_ENCRYPT,
MONGOCRYPT_KMS_GCP_DECRYPT};
bool is_retryable = false;
for (size_t i = 0; i < sizeof(retryable_types) / sizeof(retryable_types[0]); i++) {
if (retryable_types[i] == kms->req_type) {
is_retryable = true;
break;
}
}
if (!is_retryable) {
if (!_is_retryable_req(kms->req_type)) {
CLIENT_ERR("KMS request failed due to network error");
return false;
}
Expand Down
17 changes: 17 additions & 0 deletions src/mongocrypt.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1188,6 +1188,23 @@ bool mongocrypt_kms_ctx_feed(mongocrypt_kms_ctx_t *kms, mongocrypt_binary_t *byt
MONGOCRYPT_EXPORT
bool mongocrypt_kms_ctx_fail(mongocrypt_kms_ctx_t *kms);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

mongocrypt_kms_ctx_fail currently still refers to "network error" in error messages and increments the retry attempts:

ASSERT_OK(mongocrypt_kms_ctx_feed(kms, retryable_http), kms); // Increments kms->attempts.
ASSERT(mongocrypt_kms_ctx_should_retry(kms));
ASSERT_OK(mongocrypt_kms_ctx_fail(kms), kms); // Increments kms->attempts again.

I like the direction towards simplifying. Suggest resetting the KMS parser in set_retry. I expect this will reset the parser for both HTTP and network error and could further simplify the retry pattern:

while (true) {
    try {
        // TODO: Write and read from socket.
    } catch (e : NetworkError) {
        // Indicate network error:
        mongocrypt_kms_ctx_fail (kms);
    } finally {
        if (mongocrypt_kms_ctx_should_retry (kms)) {
            // Can retry due to a HTTP or network error.
            continue;
        }
    }
}

That way, mongocrypt_kms_ctx_fail does not change its meaning (still means only network error).

Testing this change in 456d523 appears to work.


/**
* Reset a KMS context allowing it to be retried.
*
* @param[in] kms The @ref mongocrypt_kms_ctx_t.
*/
MONGOCRYPT_EXPORT
void mongocrypt_kms_ctx_reset(mongocrypt_kms_ctx_t *kms);

/**
* Indicate if a KMS context is completed but should be retried due to an HTTP error.
*
* @param[in] kms The @ref mongocrypt_kms_ctx_t.
* @return A boolean indicating whether the failed request is complete but should be retried.
*/
MONGOCRYPT_EXPORT
bool mongocrypt_kms_ctx_should_retry_http(mongocrypt_kms_ctx_t *kms);

/**
* Get the status associated with a @ref mongocrypt_kms_ctx_t object.
*
Expand Down
3 changes: 3 additions & 0 deletions test/data/kms-aws/encrypt-response-partial.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
HTTP/1.1 200 OK
x-amzn-RequestId: deeb35e5-4ecb-4bf1-9af5-84a54ff0af0e
Content-Type: appli
54 changes: 54 additions & 0 deletions test/test-mongocrypt-datakey.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -427,6 +427,33 @@ static void _test_create_datakey_with_retry(_mongocrypt_tester_t *tester) {
mongocrypt_destroy(crypt);
}

// Test that an HTTP error is retried in-place.
{
mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);
ASSERT_OK(
mongocrypt_ctx_setopt_key_encryption_key(ctx,
TEST_BSON("{'provider': 'aws', 'key': 'foo', 'region': 'bar'}")),
ctx);
ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);
ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS);
mongocrypt_kms_ctx_t *kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
ASSERT_OK(kms_ctx, ctx);
// Expect no sleep is requested before any error.
ASSERT_CMPINT64(mongocrypt_kms_ctx_usleep(kms_ctx), ==, 0);
// Feed a retryable HTTP error.
ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/rmd/kms-decrypt-reply-429.txt")), kms_ctx);
// In-place retry is indicated.
ASSERT(mongocrypt_kms_ctx_should_retry_http(kms_ctx));
mongocrypt_kms_ctx_reset(kms_ctx);
// Feed a successful response.
ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-aws/encrypt-response.txt")), kms_ctx);
ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
_mongocrypt_tester_run_ctx_to(tester, ctx, MONGOCRYPT_CTX_DONE);
mongocrypt_ctx_destroy(ctx);
mongocrypt_destroy(crypt);
}

// Test that a network error is retried.
{
mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
Expand Down Expand Up @@ -454,6 +481,33 @@ static void _test_create_datakey_with_retry(_mongocrypt_tester_t *tester) {
mongocrypt_destroy(crypt);
}

// Test that a network error is retried in-place.
{
mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);
ASSERT_OK(
mongocrypt_ctx_setopt_key_encryption_key(ctx,
TEST_BSON("{'provider': 'aws', 'key': 'foo', 'region': 'bar'}")),
ctx);
ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);
ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS);
mongocrypt_kms_ctx_t *kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
ASSERT_OK(kms_ctx, ctx);
// Expect no sleep is requested before any error.
ASSERT_CMPINT64(mongocrypt_kms_ctx_usleep(kms_ctx), ==, 0);
// Feed part of a response
ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-aws/encrypt-response-partial.txt")),
kms_ctx);
// Assume a network error and reset the context.
mongocrypt_kms_ctx_reset(kms_ctx);
// Feed a successful response.
ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-aws/encrypt-response.txt")), kms_ctx);
ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
_mongocrypt_tester_run_ctx_to(tester, ctx, MONGOCRYPT_CTX_DONE);
mongocrypt_ctx_destroy(ctx);
mongocrypt_destroy(crypt);
}

// Test that an oauth request is retried for a network error.
{
mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
Expand Down
Loading