Delivery exact bytes for read window (#600)

TingDaoK · web-flow · commit a109fa6965f8 · 2025-11-24T14:50:16.000-08:00
diff --git a/include/aws/s3/private/s3_meta_request_impl.h b/include/aws/s3/private/s3_meta_request_impl.h
@@ -66,6 +66,7 @@ struct aws_s3_meta_request_event {
         /* data for AWS_S3_META_REQUEST_EVENT_RESPONSE_BODY */
         struct {
             struct aws_s3_request *completed_request;
+            size_t bytes_delivered;
         } response_body;
 
         /* data for AWS_S3_META_REQUEST_EVENT_PROGRESS */
@@ -224,12 +225,15 @@ struct aws_s3_meta_request {
 
         /* Task for delivering events on the meta-request's io_event_loop thread.
          * We do this to ensure a meta-request's callbacks are fired sequentially and non-overlapping.
-         * If `event_delivery_array` has items in it, then this task is scheduled.
+         * If `event_delivery_task_scheduled` is true, then this task is scheduled.
          * If `event_delivery_active` is true, then this task is actively running.
          * Delivery is not 100% complete until `event_delivery_array` is empty AND `event_delivery_active` is false
          * (use aws_s3_meta_request_are_events_out_for_delivery_synced()  to check) */
         struct aws_task event_delivery_task;
 
+        /* Whether or not event delivery is currently scheduled. */
+        uint32_t event_delivery_task_scheduled : 1;
+
         /* Array of `struct aws_s3_meta_request_event` to deliver when the `event_delivery_task` runs. */
         struct aws_array_list event_delivery_array;
 
@@ -292,6 +296,10 @@ struct aws_s3_meta_request {
 
         /* The range start for the next response body delivery */
         uint64_t next_deliver_range_start;
+
+        /* Total number of bytes that have been attempted to be delivered. (Will equal the sum of succeeded and
+         * failed.)*/
+        uint64_t num_bytes_delivery_completed;
     } io_threaded_data;
 
     const bool should_compute_content_md5;
diff --git a/include/aws/s3/s3_client.h b/include/aws/s3/s3_client.h
@@ -606,7 +606,8 @@ struct aws_s3_client_config {
      *
      * WARNING: This feature is experimental.
      * Currently, backpressure is only applied to GetObject requests which are split into multiple parts,
-     * and you may still receive some data after the window reaches 0.
+     * - If you set body_callback, no more data will be delivered once the window reaches 0.
+     * - If you set body_callback_ex, you may still receive some data after the window reaches 0. TODO: fix it.
      */
     bool enable_read_backpressure;
 
@@ -1220,7 +1221,8 @@ struct aws_future_void *aws_s3_meta_request_write(
  *
  * WARNING: This feature is experimental.
  * Currently, backpressure is only applied to GetObject requests which are split into multiple parts,
- * and you may still receive some data after the window reaches 0.
+ * - If you set body_callback, no more data will be delivered once the window reaches 0.
+ * - If you set body_callback_ex, you may still receive some data after the window reaches 0. TODO: fix it.
  */
 AWS_S3_API
 void aws_s3_meta_request_increment_read_window(struct aws_s3_meta_request *meta_request, uint64_t bytes);
diff --git a/source/s3_meta_request.c b/source/s3_meta_request.c
@@ -412,6 +412,8 @@ void aws_s3_meta_request_increment_read_window(struct aws_s3_meta_request *meta_
     /* Response will never approach UINT64_MAX, so do a saturating sum instead of worrying about overflow */
     meta_request->synced_data.read_window_running_total =
         aws_add_u64_saturating(bytes, meta_request->synced_data.read_window_running_total);
+    /* Try to schedule the delivery task again. */
+    aws_s3_meta_request_add_event_for_delivery_synced(meta_request, NULL);
 
     aws_s3_meta_request_unlock_synced_data(meta_request);
     /* END CRITICAL SECTION */
@@ -1882,12 +1884,13 @@ void aws_s3_meta_request_add_event_for_delivery_synced(
     const struct aws_s3_meta_request_event *event) {
 
     ASSERT_SYNCED_DATA_LOCK_HELD(meta_request);
+    if (event) {
+        aws_array_list_push_back(&meta_request->synced_data.event_delivery_array, event);
+    }
 
-    aws_array_list_push_back(&meta_request->synced_data.event_delivery_array, event);
-
-    /* If the array was empty before, schedule task to deliver all events in the array.
-     * If the array already had things in it, then the task is already scheduled and will run soon. */
-    if (aws_array_list_length(&meta_request->synced_data.event_delivery_array) == 1) {
+    /* If the event delivery task is not scheduled before, and there are more to be delivered. */
+    if (!meta_request->synced_data.event_delivery_task_scheduled &&
+        aws_array_list_length(&meta_request->synced_data.event_delivery_array) > 0) {
         aws_s3_meta_request_acquire(meta_request);
 
         aws_task_init(
@@ -1896,6 +1899,7 @@ void aws_s3_meta_request_add_event_for_delivery_synced(
             meta_request,
             "s3_meta_request_event_delivery");
         aws_event_loop_schedule_task_now(meta_request->io_event_loop, &meta_request->synced_data.event_delivery_task);
+        meta_request->synced_data.event_delivery_task_scheduled = true;
     }
 }
 
@@ -1976,24 +1980,37 @@ static void s_s3_meta_request_event_delivery_task(struct aws_task *task, void *a
     struct aws_array_list *event_delivery_array = &meta_request->io_threaded_data.event_delivery_array;
     AWS_FATAL_ASSERT(aws_array_list_length(event_delivery_array) == 0);
 
+    struct aws_array_list incomplete_deliver_events_array;
+    aws_array_list_init_dynamic(
+        &incomplete_deliver_events_array, meta_request->allocator, 1, sizeof(struct aws_s3_meta_request_event));
+
     /* If an error occurs, don't fire callbacks anymore. */
     int error_code = AWS_ERROR_SUCCESS;
     uint32_t num_parts_delivered = 0;
+    uint64_t bytes_allowed_to_deliver = 0;
+    uint64_t read_window_to_increment = 0;
 
     /* BEGIN CRITICAL SECTION */
     {
         aws_s3_meta_request_lock_synced_data(meta_request);
 
         aws_array_list_swap_contents(event_delivery_array, &meta_request->synced_data.event_delivery_array);
         meta_request->synced_data.event_delivery_active = true;
+        meta_request->synced_data.event_delivery_task_scheduled = false;
 
         if (aws_s3_meta_request_has_finish_result_synced(meta_request)) {
             error_code = AWS_ERROR_S3_CANCELED;
         }
 
+        bytes_allowed_to_deliver = meta_request->synced_data.read_window_running_total -
+                                   meta_request->io_threaded_data.num_bytes_delivery_completed;
+
         aws_s3_meta_request_unlock_synced_data(meta_request);
     }
     /* END CRITICAL SECTION */
+    if (bytes_allowed_to_deliver > SIZE_MAX) {
+        bytes_allowed_to_deliver = SIZE_MAX;
+    }
 
     /* Deliver all events */
     for (size_t event_i = 0; event_i < aws_array_list_length(event_delivery_array); ++event_i) {
@@ -2003,15 +2020,44 @@ static void s_s3_meta_request_event_delivery_task(struct aws_task *task, void *a
 
             case AWS_S3_META_REQUEST_EVENT_RESPONSE_BODY: {
                 struct aws_s3_request *request = event.u.response_body.completed_request;
+                size_t bytes_delivered_for_request = event.u.response_body.bytes_delivered;
                 AWS_ASSERT(meta_request == request->meta_request);
+                bool delivery_incomplete = false;
                 struct aws_byte_cursor response_body = aws_byte_cursor_from_buf(&request->send_data.response_body);
+                if (response_body.len == 0) {
+                    /* Nothing to delivery, finish this delivery event and break out. */
+                    aws_atomic_fetch_sub(&client->stats.num_requests_streaming_response, 1);
+
+                    ++num_parts_delivered;
+                    request->send_data.metrics =
+                        s_s3_request_finish_up_and_release_metrics(request->send_data.metrics, meta_request);
+
+                    aws_s3_request_release(request);
+                    break;
+                }
+
+                if (meta_request->body_callback && meta_request->client->enable_read_backpressure) {
+                    /* If customer set the body callback, make sure we are not delivery them more than asked via the
+                     * callback. */
+                    aws_byte_cursor_advance(&response_body, bytes_delivered_for_request);
+                    if (response_body.len > (size_t)bytes_allowed_to_deliver) {
+                        response_body.len = (size_t)bytes_allowed_to_deliver;
+                        delivery_incomplete = true;
+                    }
+                    /* Update the remaining bytes we allow to delivery. */
+                    bytes_allowed_to_deliver -= response_body.len;
+                } else {
+                    /* We should not have any incomplete delivery in this case. */
+                    AWS_FATAL_ASSERT(bytes_delivered_for_request == 0);
+                }
+                uint64_t delivery_range_start = request->part_range_start + bytes_delivered_for_request;
 
                 AWS_ASSERT(request->part_number >= 1);
                 if (request->part_number == 1) {
-                    meta_request->io_threaded_data.next_deliver_range_start = request->part_range_start;
+                    meta_request->io_threaded_data.next_deliver_range_start = delivery_range_start;
                 }
                 /* Make sure the response body is delivered in the sequential order */
-                AWS_FATAL_ASSERT(request->part_range_start == meta_request->io_threaded_data.next_deliver_range_start);
+                AWS_FATAL_ASSERT(delivery_range_start == meta_request->io_threaded_data.next_deliver_range_start);
                 meta_request->io_threaded_data.next_deliver_range_start += response_body.len;
 
                 if (error_code == AWS_ERROR_SUCCESS && response_body.len > 0) {
@@ -2047,15 +2093,15 @@ static void s_s3_meta_request_event_delivery_task(struct aws_task *task, void *a
                                     aws_error_name(error_code));
                             }
                             if (meta_request->client->enable_read_backpressure) {
-                                aws_s3_meta_request_increment_read_window(meta_request, response_body.len);
+                                read_window_to_increment += response_body.len;
                             }
                         } else if (
                             meta_request->body_callback_ex != NULL &&
                             meta_request->body_callback_ex(
                                 meta_request,
                                 &response_body,
                                 (struct aws_s3_meta_request_receive_body_extra_info){
-                                    .range_start = request->part_range_start, .ticket = request->ticket},
+                                    .range_start = delivery_range_start, .ticket = request->ticket},
                                 meta_request->user_data)) {
                             error_code = aws_last_error_or_unknown();
                             AWS_LOGF_ERROR(
@@ -2067,7 +2113,7 @@ static void s_s3_meta_request_event_delivery_task(struct aws_task *task, void *a
                         } else if (
                             meta_request->body_callback != NULL &&
                             meta_request->body_callback(
-                                meta_request, &response_body, request->part_range_start, meta_request->user_data)) {
+                                meta_request, &response_body, delivery_range_start, meta_request->user_data)) {
 
                             error_code = aws_last_error_or_unknown();
                             AWS_LOGF_ERROR(
@@ -2087,13 +2133,25 @@ static void s_s3_meta_request_event_delivery_task(struct aws_task *task, void *a
                         }
                     }
                 }
-                aws_atomic_fetch_sub(&client->stats.num_requests_streaming_response, 1);
+                event.u.response_body.bytes_delivered += response_body.len;
+                meta_request->io_threaded_data.num_bytes_delivery_completed += response_body.len;
 
-                ++num_parts_delivered;
-                request->send_data.metrics =
-                    s_s3_request_finish_up_and_release_metrics(request->send_data.metrics, meta_request);
+                if (!delivery_incomplete || error_code != AWS_ERROR_SUCCESS) {
+                    /* We completed the delivery for this request. */
+                    aws_atomic_fetch_sub(&client->stats.num_requests_streaming_response, 1);
 
-                aws_s3_request_release(request);
+                    ++num_parts_delivered;
+                    request->send_data.metrics =
+                        s_s3_request_finish_up_and_release_metrics(request->send_data.metrics, meta_request);
+
+                    aws_s3_request_release(request);
+                } else {
+                    /* We didn't complete the delivery for this request and no error happened */
+                    /* Push to the front of the queue and wait for the next tick to deliver the rest of the bytes. */
+                    /* Note: we push to the front of the array since when we move those incomplete events back to the
+                     * synced_queue, we need to make sure it still has the same order. */
+                    aws_array_list_push_front(&incomplete_deliver_events_array, &event);
+                }
             } break;
 
             case AWS_S3_META_REQUEST_EVENT_PROGRESS: {
@@ -2149,12 +2207,21 @@ static void s_s3_meta_request_event_delivery_task(struct aws_task *task, void *a
         if (error_code != AWS_ERROR_SUCCESS) {
             aws_s3_meta_request_set_fail_synced(meta_request, NULL, error_code);
         }
+        /* Push the iocompileted events back to the queue */
+        for (size_t i = 0; i < aws_array_list_length(&incomplete_deliver_events_array); ++i) {
+            struct aws_s3_meta_request_event event;
+            aws_array_list_get_at(&incomplete_deliver_events_array, &event, i);
+            /* Push the incomplete one to the front of the queue. */
+            aws_array_list_push_front(&meta_request->synced_data.event_delivery_array, &event);
+        }
 
         meta_request->synced_data.num_parts_delivery_completed += num_parts_delivered;
         meta_request->synced_data.event_delivery_active = false;
         aws_s3_meta_request_unlock_synced_data(meta_request);
     }
     /* END CRITICAL SECTION */
+    aws_s3_meta_request_increment_read_window(meta_request, read_window_to_increment);
+    aws_array_list_clean_up(&incomplete_deliver_events_array);
 
     aws_s3_client_schedule_process_work(client);
     aws_s3_meta_request_release(meta_request);
diff --git a/tests/s3_data_plane_tests.c b/tests/s3_data_plane_tests.c
@@ -1726,6 +1726,7 @@ static int s_apply_backpressure_until_meta_request_finish(
     size_t part_size,
     size_t window_initial_size,
     uint64_t window_increment_size) {
+    (void)part_size;
 
     /* Remember the last time something happened (we received download data, or incremented read window) */
     uint64_t last_time_something_happened;
@@ -1753,12 +1754,13 @@ static int s_apply_backpressure_until_meta_request_finish(
         size_t received_body_size_delta = aws_atomic_exchange_int(&test_results->received_body_size_delta, 0);
         accumulated_data_size += (uint64_t)received_body_size_delta;
 
-        /* Check that we haven't received more data than the window allows.
-         * TODO: Stop allowing "hacky wiggle room". The current implementation
-         *       may push more bytes to the user (up to 1 part) than they've asked for. */
-        uint64_t hacky_wiggle_room = part_size;
-        uint64_t max_data_allowed = accumulated_window_increments + hacky_wiggle_room;
-        ASSERT_TRUE(accumulated_data_size <= max_data_allowed, "Received more data than the read window allows");
+        /* Check that we haven't received more data than the window allows */
+        uint64_t max_data_allowed = accumulated_window_increments;
+        ASSERT_TRUE(
+            accumulated_data_size <= max_data_allowed,
+            "Received more data than the read window allows accumulated_data_size: %zu, max_data_allowed: %zu",
+            (size_t)accumulated_data_size,
+            (size_t)max_data_allowed);
 
         /* If we're done, we're done */
         if (done) {
@@ -1891,7 +1893,7 @@ static int s_test_s3_get_object_backpressure_small_increments(struct aws_allocat
     size_t file_size = 1 * 1024 * 1024; /* Test downloads 1MB file */
     size_t part_size = file_size / 4;
     size_t window_initial_size = 1024;
-    uint64_t window_increment_size = part_size / 2;
+    uint64_t window_increment_size = part_size / 4;
     return s_test_s3_get_object_backpressure_helper(
         allocator, part_size, window_initial_size, window_increment_size, false);
 }
