try to fix the flaky test

Author: TingDaoK

include/aws/s3/private/s3_client_impl.h

@@ -178,7 +178,8 @@ struct aws_s3_client_vtable {
         struct aws_http_connection *client_connection,
         const struct aws_http_make_request_options *options);
 
-    void (*after_prepare_upload_part_finish)(struct aws_s3_request *request, struct aws_http_message *message);
+    /********************* TEST ONLY STUB **************************/
+    void (*after_prepare_upload_part_finish_stub)(struct aws_s3_request *request, struct aws_http_message *message);
 };
 
 struct aws_s3_upload_part_timeout_stats {

include/aws/s3/private/s3_meta_request_impl.h

@@ -122,6 +122,10 @@ struct aws_s3_meta_request_vtable {
 
     /* Pause the given request */
     int (*pause)(struct aws_s3_meta_request *meta_request, struct aws_s3_meta_request_resume_token **resume_token);
+
+    /********************* TEST ONLY STUB **************************/
+    /* A stub to the update implementation from meta request with the lock held. Only for tests. */
+    bool (*synced_update_stub)(struct aws_s3_meta_request *meta_request);
 };
 
 /**

source/s3_auto_ranged_get.c

@@ -202,6 +202,11 @@ static bool s_s3_auto_ranged_get_update(
     /* BEGIN CRITICAL SECTION */
     {
         aws_s3_meta_request_lock_synced_data(meta_request);
+        if (meta_request->vtable->synced_update_stub && meta_request->vtable->synced_update_stub(meta_request)) {
+            /* TEST ONLY, allow test to stub here. */
+            aws_s3_meta_request_unlock_synced_data(meta_request);
+            return true;
+        }
 
         /* If nothing has set the "finish result" then this meta request is still in progress, and we can potentially
          * send additional requests. */

source/s3_auto_ranged_put.c

@@ -486,6 +486,11 @@ static bool s_s3_auto_ranged_put_update(
     /* BEGIN CRITICAL SECTION */
     {
         aws_s3_meta_request_lock_synced_data(meta_request);
+        if (meta_request->vtable->synced_update_stub && meta_request->vtable->synced_update_stub(meta_request)) {
+            /* TEST ONLY, allow test to stub here. */
+            aws_s3_meta_request_unlock_synced_data(meta_request);
+            return true;
+        }
 
         if (!aws_s3_meta_request_has_finish_result_synced(meta_request)) {
             /* If resuming and list part has not been sent, do it now. */
@@ -1248,9 +1253,9 @@ static void s_s3_prepare_upload_part_finish(struct aws_s3_prepare_upload_part_jo
         aws_future_http_message_set_error(part_prep->on_complete, aws_last_error());
         goto on_done;
     }
-    if (client->vtable->after_prepare_upload_part_finish) {
+    if (client->vtable->after_prepare_upload_part_finish_stub) {
         /* TEST ONLY, allow test to stub here. */
-        client->vtable->after_prepare_upload_part_finish(request, message);
+        client->vtable->after_prepare_upload_part_finish_stub(request, message);
     }
 
     /* Success! */
@@ -1779,6 +1784,8 @@ static void s_s3_auto_ranged_put_request_finished(
     aws_s3_meta_request_unlock_synced_data(meta_request);
 }
 
+/* NOTES: the implementation has been copy/pasted to `s_pause_meta_request_synced` for testing purpose, if changed made
+ * here, please update the other function correspondingly. */
 static int s_s3_auto_ranged_put_pause(
     struct aws_s3_meta_request *meta_request,
     struct aws_s3_meta_request_resume_token **out_resume_token) {

source/s3_default_meta_request.c

@@ -160,7 +160,11 @@ static bool s_s3_meta_request_default_update(
     /* BEGIN CRITICAL SECTION */
     {
         aws_s3_meta_request_lock_synced_data(meta_request);
-
+        if (meta_request->vtable->synced_update_stub && meta_request->vtable->synced_update_stub(meta_request)) {
+            /* TEST ONLY, allow test to stub here. */
+            aws_s3_meta_request_unlock_synced_data(meta_request);
+            return true;
+        }
         if (!aws_s3_meta_request_has_finish_result_synced(meta_request)) {
 
             /* If the request hasn't been sent, then create and send it now. */

tests/s3_cancel_tests.c

@@ -40,12 +40,48 @@ struct s3_cancel_test_user_data {
     bool abort_successful;
 };
 
-static bool s_s3_meta_request_update_cancel_test(
+/* TODO: this is an unfortunately copy/paste from `s_s3_auto_ranged_put_pause`, we need to hold the lock to invoke the
+ * pause for the test to avoid the checks changed after checking it and used afterward. In between move the
+ * synced_section from s_s3_auto_ranged_put_pause to a function that can be used here and copy/paste it here. I picked
+ * copy/paste it to avoid exposing a weird test only function to our API.  */
+static void s_pause_meta_request_synced(
     struct aws_s3_meta_request *meta_request,
-    uint32_t flags,
-    struct aws_s3_request **out_request) {
+    struct aws_s3_meta_request_resume_token **out_resume_token) {
+
+    struct aws_s3_auto_ranged_put *auto_ranged_put = meta_request->impl;
+    /* upload can be in one of several states:
+     * - not started, i.e. we didn't even call crete mpu yet - return success,
+     *   token is NULL and cancel the upload
+     * - in the middle of upload - return success, create token and cancel
+     *     upload
+     * - complete MPU started - return success, generate token and try to cancel
+     *   complete MPU
+     */
+    if (auto_ranged_put->synced_data.create_multipart_upload_completed) {
+
+        *out_resume_token = aws_s3_meta_request_resume_token_new(meta_request->allocator);
+
+        (*out_resume_token)->type = AWS_S3_META_REQUEST_TYPE_PUT_OBJECT;
+        (*out_resume_token)->multipart_upload_id =
+            aws_string_clone_or_reuse(meta_request->allocator, auto_ranged_put->upload_id);
+        (*out_resume_token)->part_size = meta_request->part_size;
+        (*out_resume_token)->total_num_parts = auto_ranged_put->total_num_parts_from_content_length;
+        (*out_resume_token)->num_parts_completed = auto_ranged_put->synced_data.num_parts_completed;
+    }
+
+    /**
+     * Cancels the meta request using the PAUSED flag to avoid deletion of uploaded parts.
+     * This allows the client to resume the upload later, setting the persistable state in the meta request options.
+     */
+    aws_s3_meta_request_set_fail_synced(meta_request, NULL, AWS_ERROR_S3_PAUSED);
+
+    aws_s3_meta_request_cancel_cancellable_requests_synced(meta_request, AWS_ERROR_S3_PAUSED);
+    aws_s3_meta_request_cancel_pending_buffer_futures_synced(meta_request, AWS_ERROR_S3_PAUSED);
+}
+
+static bool s_s3_meta_request_cancel_test_synced_update_stub(struct aws_s3_meta_request *meta_request) {
     AWS_PRECONDITION(meta_request);
-    AWS_PRECONDITION(out_request);
+    ASSERT_SYNCED_DATA_LOCK_HELD(meta_request);
 
     struct aws_s3_meta_request_test_results *results = meta_request->user_data;
     struct aws_s3_tester *tester = results->tester;
@@ -57,8 +93,6 @@ static bool s_s3_meta_request_update_cancel_test(
     bool call_cancel_or_pause = false;
     bool block_update = false;
 
-    aws_s3_meta_request_lock_synced_data(meta_request);
-
     switch (cancel_test_user_data->type) {
         case S3_UPDATE_CANCEL_TYPE_NO_CANCEL:
             break;
@@ -131,23 +165,20 @@ static bool s_s3_meta_request_update_cancel_test(
             break;
     }
 
-    aws_s3_meta_request_unlock_synced_data(meta_request);
     if (call_cancel_or_pause) {
         if (cancel_test_user_data->pause) {
-            aws_s3_meta_request_pause(meta_request, &cancel_test_user_data->resume_token);
+            s_pause_meta_request_synced(meta_request, &cancel_test_user_data->resume_token);
         } else {
-            aws_s3_meta_request_cancel(meta_request);
+            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);
+            aws_s3_meta_request_cancel_pending_buffer_futures_synced(meta_request, AWS_ERROR_S3_CANCELED);
         }
     }
 
     if (block_update) {
         return true;
     }
-
-    struct aws_s3_meta_request_vtable *original_meta_request_vtable =
-        aws_s3_tester_get_meta_request_vtable_patch(tester, 0)->original_vtable;
-
-    return original_meta_request_vtable->update(meta_request, flags, out_request);
+    return false;
 }
 
 static void s_s3_meta_request_finished_request_cancel_test(
@@ -188,7 +219,7 @@ static struct aws_s3_meta_request *s_meta_request_factory_patch_update_cancel_te
 
     struct aws_s3_meta_request_vtable *patched_meta_request_vtable =
         aws_s3_tester_patch_meta_request_vtable(tester, meta_request, NULL);
-    patched_meta_request_vtable->update = s_s3_meta_request_update_cancel_test;
+    patched_meta_request_vtable->synced_update_stub = s_s3_meta_request_cancel_test_synced_update_stub;
     patched_meta_request_vtable->finished_request = s_s3_meta_request_finished_request_cancel_test;
 
     return meta_request;

tests/s3_mock_server_tests.c

@@ -487,7 +487,7 @@ TEST_CASE(multipart_upload_checksum_with_retry_before_finish_mock_server) {
     struct aws_s3_client *client = NULL;
     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->after_prepare_upload_part_finish =
+    patched_client_vtable->after_prepare_upload_part_finish_stub =
         s_after_prepare_upload_part_finish_retry_before_finish_sending;
 
     struct aws_byte_cursor object_path = aws_byte_cursor_from_c_str("/default");
@@ -547,7 +547,7 @@ TEST_CASE(multipart_upload_checksum_with_retry_mock_server) {
     struct aws_s3_client *client = NULL;
     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->after_prepare_upload_part_finish = s_after_prepare_upload_part_finish;
+    patched_client_vtable->after_prepare_upload_part_finish_stub = s_after_prepare_upload_part_finish;
 
     struct aws_byte_cursor object_path = aws_byte_cursor_from_c_str("/default");
     {
@@ -629,7 +629,7 @@ TEST_CASE(multipart_upload_checksum_fio_with_retry_mock_server) {
     struct aws_s3_client *client = NULL;
     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->after_prepare_upload_part_finish =
+    patched_client_vtable->after_prepare_upload_part_finish_stub =
         s_after_prepare_upload_part_finish_retry_before_finish_sending;
     struct aws_s3_file_io_options fio_opts = {
         .should_stream = true,