meta request: support for should_continue cancellation callback

include/aws/s3/private/s3_meta_request_impl.h

@@ -169,6 +169,7 @@ struct aws_s3_meta_request {
 
     /* Customer specified callbacks. */
     aws_s3_meta_request_headers_callback_fn *headers_callback;
+    aws_s3_meta_request_should_continue_fn *continue_callback;
     aws_s3_meta_request_receive_body_callback_fn *body_callback;
     aws_s3_meta_request_finish_fn *finish_callback;
     aws_s3_meta_request_shutdown_fn *shutdown_callback;

include/aws/s3/s3_client.h

@@ -178,6 +178,11 @@ typedef int(aws_s3_meta_request_receive_body_callback_fn)(
     /* User data specified by aws_s3_meta_request_options.*/
     void *user_data);
 
+/**
+ * Invoked to poll whether the user has canceled the meta request.
+ */
+typedef bool(aws_s3_meta_request_should_continue_fn)(void *user_data);
+
 /**
  * Invoked when the entire meta request execution is complete.
  */
@@ -767,6 +772,12 @@ struct aws_s3_meta_request_options {
      */
     aws_s3_meta_request_receive_body_callback_fn *body_callback;
 
+    /**
+     * Invoked to check whether the user has set the cancellation flag for this request.
+     * See `aws_s3_meta_request_should_continue_fn`.
+     */
+    aws_s3_meta_request_should_continue_fn *continue_callback;
+
     /**
      * Invoked when the entire meta request execution is complete.
      * See `aws_s3_meta_request_finish_fn`.

source/s3_meta_request.c

@@ -306,6 +306,7 @@ int aws_s3_meta_request_init_base(
 
     meta_request->meta_request_level_running_response_sum = NULL;
     meta_request->user_data = options->user_data;
+    meta_request->continue_callback = options->continue_callback;
     meta_request->progress_callback = options->progress_callback;
     meta_request->telemetry_callback = options->telemetry_callback;
     meta_request->upload_review_callback = options->upload_review_callback;
@@ -452,10 +453,19 @@ bool aws_s3_meta_request_has_finish_result_synced(struct aws_s3_meta_request *me
     ASSERT_SYNCED_DATA_LOCK_HELD(meta_request);
 
     if (!meta_request->synced_data.finish_result_set) {
-        return false;
+        /*
+         * Perform the equivalent of calling aws_s3_meta_request_cancel() if continue_callback() indicates
+         * that the request should be canceled. This is a work-around to enable canceling ongoing requests
+         * via a user-provided function. The work-around is needed until the AWS C++ SDK supports Cancel().
+         * This function was chosen since it is evaluated several times during the lifetime of a request.
+         */
+        if (meta_request->continue_callback && !meta_request->continue_callback(meta_request->user_data)) {
+            aws_s3_meta_request_set_fail_synced(meta_request, NULL, AWS_ERROR_S3_CANCELED);
+            aws_s3_meta_request_cancel_cancellable_requests_synced(meta_request, AWS_ERROR_S3_CANCELED);
+        }
     }
 
-    return true;
+    return meta_request->synced_data.finish_result_set;
 }
 
 struct aws_s3_meta_request *aws_s3_meta_request_acquire(struct aws_s3_meta_request *meta_request) {

tests/CMakeLists.txt

@@ -64,6 +64,8 @@ add_net_test_case(test_s3_cancel_mpd_empty_object_get_with_part_number_1_sent)
 add_net_test_case(test_s3_cancel_mpd_empty_object_get_with_part_number_1_completed)
 add_net_test_case(test_s3_cancel_mpd_pending_streaming)
 add_net_test_case(test_s3_cancel_prepare)
+add_net_test_case(test_s3_continue_callback_ok)
+add_net_test_case(test_s3_continue_callback_cancels)
 
 add_net_test_case(test_s3_get_object_tls_disabled)
 add_net_test_case(test_s3_get_object_tls_enabled)

tests/s3_cancel_tests.c

@@ -705,6 +705,11 @@ static void s_test_s3_cancel_prepare_meta_request_prepare_request_on_original_do
 
     ++test_user_data->request_prepare_counters[request->request_tag];
 
+    if (results->continue_callback != NULL) {
+        /* Cancel via s_test_continue_callback_cancels_after_first_part(). */
+        return AWS_OP_SUCCESS;
+    }
+
     /* Cancel after the first part is prepared, preventing any additional parts from being prepared. */
     if (request->request_tag == AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_PART &&
         test_user_data->request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_PART] == 1) {
@@ -787,3 +792,112 @@ static int s_test_s3_cancel_prepare(struct aws_allocator *allocator, void *ctx)
 
     return 0;
 }
+
+static bool s_test_continue_callback_cancels_after_first_part(void *user_data) {
+    struct aws_s3_meta_request_test_results *results = user_data;
+    AWS_ASSERT(results != NULL);
+
+    struct aws_s3_tester *tester = results->tester;
+    AWS_ASSERT(tester != NULL);
+
+    /* Similar to s_test_s3_cancel_prepare_meta_request_prepare_request: cancel after the first part. */
+    struct test_s3_cancel_prepare_user_data *test_user_data = tester->user_data;
+    return test_user_data->request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_PART] == 0;
+}
+
+AWS_TEST_CASE(test_s3_continue_callback_cancels, s_test_s3_continue_callback_cancels)
+static int s_test_s3_continue_callback_cancels(struct aws_allocator *allocator, void *ctx) {
+    (void)ctx;
+
+    struct aws_s3_tester tester;
+    ASSERT_SUCCESS(aws_s3_tester_init(allocator, &tester));
+
+    struct test_s3_cancel_prepare_user_data test_user_data;
+    AWS_ZERO_STRUCT(test_user_data);
+    tester.user_data = &test_user_data;
+
+    struct aws_s3_client *client = NULL;
+
+    struct aws_s3_tester_client_options client_options;
+    AWS_ZERO_STRUCT(client_options);
+
+    ASSERT_SUCCESS(aws_s3_tester_client_new(&tester, &client_options, &client));
+
+    struct aws_s3_client_vtable *patched_client_vtable = aws_s3_tester_patch_client_vtable(&tester, client, NULL);
+    patched_client_vtable->meta_request_factory = s_test_s3_cancel_prepare_meta_request_factory;
+
+    {
+        struct aws_s3_tester_meta_request_options options = {
+            .allocator = allocator,
+            .client = client,
+            .meta_request_type = AWS_S3_META_REQUEST_TYPE_PUT_OBJECT,
+            .validate_type = AWS_S3_TESTER_VALIDATE_TYPE_EXPECT_FAILURE,
+            .continue_callback = s_test_continue_callback_cancels_after_first_part,
+            .put_options =
+                {
+                    .ensure_multipart = true,
+                },
+        };
+        struct aws_s3_meta_request_test_results meta_request_test_results;
+        AWS_ZERO_STRUCT(meta_request_test_results);
+
+        ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &options, &meta_request_test_results));
+        ASSERT_TRUE(meta_request_test_results.finished_error_code == AWS_ERROR_S3_CANCELED);
+
+        aws_s3_meta_request_test_results_clean_up(&meta_request_test_results);
+    }
+
+    ASSERT_TRUE(
+        test_user_data.request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_CREATE_MULTIPART_UPLOAD] == 1);
+    ASSERT_TRUE(test_user_data.request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_PART] == 1);
+    ASSERT_TRUE(
+        test_user_data.request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_ABORT_MULTIPART_UPLOAD] == 1);
+    ASSERT_TRUE(
+        test_user_data.request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_COMPLETE_MULTIPART_UPLOAD] == 0);
+
+    aws_s3_client_release(client);
+    aws_s3_tester_clean_up(&tester);
+
+    return 0;
+}
+
+static bool s_test_continue_callback_ok(void *user_data) {
+    (void)user_data;
+    return true;
+}
+
+AWS_TEST_CASE(test_s3_continue_callback_ok, s_test_s3_continue_callback_ok)
+static int s_test_s3_continue_callback_ok(struct aws_allocator *allocator, void *ctx) {
+    (void)ctx;
+
+    struct aws_s3_tester tester;
+    ASSERT_SUCCESS(aws_s3_tester_init(allocator, &tester));
+
+    struct aws_s3_client *client = NULL;
+
+    struct aws_s3_tester_client_options client_options;
+    AWS_ZERO_STRUCT(client_options);
+
+    ASSERT_SUCCESS(aws_s3_tester_client_new(&tester, &client_options, &client));
+
+    {
+        struct aws_s3_tester_meta_request_options options = {
+            .allocator = allocator,
+            .client = client,
+            .meta_request_type = AWS_S3_META_REQUEST_TYPE_PUT_OBJECT,
+            .validate_type = AWS_S3_TESTER_VALIDATE_TYPE_EXPECT_SUCCESS,
+            .continue_callback = s_test_continue_callback_ok,
+            .put_options =
+                {
+                    .ensure_multipart = true,
+                },
+        };
+
+        ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &options, NULL));
+    }
+
+    aws_s3_client_release(client);
+    aws_s3_tester_clean_up(&tester);
+
+    return 0;
+}

tests/s3_tester.c

@@ -114,6 +114,16 @@ static int s_s3_test_meta_request_header_callback(
     return AWS_OP_SUCCESS;
 }
 
+static bool s_s3_test_meta_request_continue_callback(void *user_data) {
+    struct aws_s3_meta_request_test_results *meta_request_test_results =
+        (struct aws_s3_meta_request_test_results *)user_data;
+
+    if (meta_request_test_results->continue_callback != NULL) {
+        return meta_request_test_results->continue_callback(user_data);
+    }
+    return true;
+}
+
 static int s_s3_test_meta_request_body_callback(
     struct aws_s3_meta_request *meta_request,
     const struct aws_byte_cursor *body,
@@ -506,6 +516,9 @@ int aws_s3_tester_bind_meta_request(
     ASSERT_TRUE(options->headers_callback == NULL);
     options->headers_callback = s_s3_test_meta_request_header_callback;
 
+    ASSERT_TRUE(options->continue_callback == NULL);
+    options->continue_callback = s_s3_test_meta_request_continue_callback;
+
     ASSERT_TRUE(options->body_callback == NULL);
     options->body_callback = s_s3_test_meta_request_body_callback;
 
@@ -1732,6 +1745,7 @@ int aws_s3_tester_send_meta_request_with_options(
     }
 
     out_results->headers_callback = options->headers_callback;
+    out_results->continue_callback = options->continue_callback;
     out_results->body_callback = options->body_callback;
     out_results->finish_callback = options->finish_callback;
     out_results->progress_callback = options->progress_callback;

tests/s3_tester.h

@@ -172,6 +172,7 @@ struct aws_s3_tester_meta_request_options {
     bool use_s3express_signing;
 
     aws_s3_meta_request_headers_callback_fn *headers_callback;
+    aws_s3_meta_request_should_continue_fn *continue_callback;
     aws_s3_meta_request_receive_body_callback_fn *body_callback;
     aws_s3_meta_request_finish_fn *finish_callback;
     aws_s3_meta_request_progress_fn *progress_callback;
@@ -238,6 +239,7 @@ struct aws_s3_meta_request_test_results {
     struct aws_s3_tester *tester;
 
     aws_s3_meta_request_headers_callback_fn *headers_callback;
+    aws_s3_meta_request_should_continue_fn *continue_callback;
     aws_s3_meta_request_receive_body_callback_fn *body_callback;
     aws_s3_meta_request_finish_fn *finish_callback;
     aws_s3_meta_request_progress_fn *progress_callback;