collect part range for auto-range-put as well (#588)

Author: TingDaoK

include/aws/s3/private/s3_request.h

@@ -241,11 +241,9 @@ struct aws_s3_request {
     struct aws_s3_buffer_ticket *ticket;
 
     /* Beginning range of this part. */
-    /* TODO currently only used by auto_range_get, could be hooked up to auto_range_put as well. */
     uint64_t part_range_start;
 
     /* Last byte of this part.*/
-    /* TODO currently only used by auto_range_get, could be hooked up to auto_range_put as well. */
     uint64_t part_range_end;
 
     /* Part number that this request refers to.  If this is not a part, this can be 0.  (S3 Part Numbers start at 1.)

source/s3_auto_ranged_put.c

@@ -759,31 +759,33 @@ static bool s_s3_auto_ranged_put_update(
 }
 
 /**
- * Helper to compute request body size.
- * Basically returns either part size or if content is not equally divisible into parts, the size of the remaining last
- * part.
+ * Helper to initialize the request ranges and content-length
+ * based on the request->part_number and meta_request->part_size
  */
-static size_t s_compute_request_body_size(
-    const struct aws_s3_meta_request *meta_request,
-    uint32_t part_number,
-    uint64_t *offset_out) {
-    AWS_PRECONDITION(meta_request);
+static int s_compute_request_body_size(const struct aws_s3_meta_request *meta_request, struct aws_s3_request *request) {
+    AWS_ERROR_PRECONDITION(meta_request);
+    AWS_ERROR_PRECONDITION(request);
 
     const struct aws_s3_auto_ranged_put *auto_ranged_put = meta_request->impl;
 
     size_t request_body_size = meta_request->part_size;
     /* Last part--adjust size to match remaining content length. */
-    if (auto_ranged_put->has_content_length && part_number == auto_ranged_put->total_num_parts_from_content_length) {
+    if (auto_ranged_put->has_content_length &&
+        request->part_number == auto_ranged_put->total_num_parts_from_content_length) {
         size_t content_remainder = (size_t)(auto_ranged_put->content_length % (uint64_t)meta_request->part_size);
 
         if (content_remainder > 0) {
             request_body_size = content_remainder;
         }
     }
-    /* The part_number starts at 1 */
-    *offset_out = (part_number - 1) * meta_request->part_size;
-
-    return request_body_size;
+    if (aws_mul_u64_checked(request->part_number - 1, meta_request->part_size, &request->part_range_start)) {
+        return AWS_OP_ERR;
+    }
+    if (aws_add_u64_checked(request->part_range_start, request_body_size - 1, &request->part_range_end)) {
+        return AWS_OP_ERR;
+    }
+    request->content_length = request_body_size;
+    return AWS_OP_SUCCESS;
 }
 
 static int s_verify_part_matches_checksum(
@@ -1085,23 +1087,24 @@ struct aws_future_http_message *s_s3_prepare_upload_part(struct aws_s3_request *
     part_prep->allocator = allocator;
     part_prep->request = request;
     part_prep->on_complete = aws_future_http_message_acquire(message_future);
+    if (s_compute_request_body_size(meta_request, request)) {
+        s_s3_prepare_upload_part_finish(part_prep, aws_last_error_or_unknown());
+        return message_future;
+    }
 
     if (request->fio_streaming) {
         /* Create the request body stream for the HTTP to read directly from the file. If retry happens, just recreate
          * it. */
         aws_input_stream_release(request->request_body_stream);
-        uint64_t offset = 0;
-        size_t request_body_size = s_compute_request_body_size(meta_request, request->part_number, &offset);
         AWS_ASSERT(meta_request->request_body_parallel_stream != NULL);
+        int error_code = AWS_ERROR_SUCCESS;
         request->request_body_stream = aws_part_streaming_input_stream_new(
             allocator,
             meta_request->request_body_parallel_stream,
             request->ticket,
-            offset,
-            request_body_size,
+            request->part_range_start,
+            (size_t)request->content_length,
             meta_request->fio_opts.direct_io);
-        request->content_length = request_body_size;
-        int error_code = AWS_ERROR_SUCCESS;
 
         if (request->num_times_prepared == 0) {
             if (s_s3_new_upload_part_info_after_body(request, meta_request, false)) {
@@ -1117,15 +1120,14 @@ struct aws_future_http_message *s_s3_prepare_upload_part(struct aws_s3_request *
          * from an upload that had been paused) */
 
         /* Read the body */
-        uint64_t offset = 0;
-        size_t request_body_size = s_compute_request_body_size(meta_request, request->part_number, &offset);
         if (request->request_body.capacity == 0) {
             AWS_FATAL_ASSERT(request->ticket);
             request->request_body = aws_s3_buffer_ticket_claim(request->ticket);
-            request->request_body.capacity = request_body_size;
+            request->request_body.capacity = (size_t)request->content_length;
         }
 
-        part_prep->asyncstep_read_part = aws_s3_meta_request_read_body(meta_request, offset, &request->request_body);
+        part_prep->asyncstep_read_part =
+            aws_s3_meta_request_read_body(meta_request, request->part_range_start, &request->request_body);
         aws_future_bool_register_callback(
             part_prep->asyncstep_read_part, s_s3_prepare_upload_part_on_read_done, part_prep);
     } else {
@@ -1159,21 +1161,21 @@ static void s_s3_prepare_upload_part_on_read_done(void *user_data) {
             request->request_body.capacity);
         goto on_done;
     }
-    /* Reading succeeded. */
-    request->content_length = request->request_body.len;
     bool is_body_stream_at_end = aws_future_bool_get_result(part_prep->asyncstep_read_part);
-
-    uint64_t offset = 0;
-    size_t request_body_size = s_compute_request_body_size(meta_request, request->part_number, &offset);
     /* If Content-Length is defined, check that we read the expected amount */
-    if (has_content_length && (request->request_body.len < request_body_size)) {
+    if (has_content_length && (request->request_body.len < request->content_length)) {
         error_code = AWS_ERROR_S3_INCORRECT_CONTENT_LENGTH;
         AWS_LOGF_ERROR(
             AWS_LS_S3_META_REQUEST,
             "id=%p: Request body is smaller than 'Content-Length' header said it would be",
             (void *)meta_request);
         goto on_done;
     }
+    if (!has_content_length && (request->request_body.len < request->content_length)) {
+        /* The stream is not ready to provide enough data as expected. Update the tracking numbers. */
+        request->content_length = request->request_body.len;
+        request->part_range_end = request->part_range_start + request->content_length - 1;
+    }
     request->is_noop = request->part_number >
                            1 && /* allow first part to have 0 length to support empty unknown content length objects. */
                        request->request_body.len == 0;

tests/s3_data_plane_tests.c

@@ -4987,6 +4987,33 @@ static int s_test_s3_round_trip_mpu_with_filepath_streaming_full_object_checksum
         AWS_TEST_FOC_CALLBACK);
 }
 
+static int s_check_metrics_upload_part_helper(
+    struct aws_s3_meta_request_test_results *test_results,
+    size_t expected_part_size,
+    size_t index) {
+    struct aws_s3_request_metrics *metrics = NULL;
+    /* The upload part can be completed out of order. Just make sure the range matches the part number */
+    ASSERT_SUCCESS(aws_array_list_get_at(&test_results->synced_data.succeed_metrics, (void **)&metrics, index));
+    enum aws_s3_request_type out_request_type = AWS_S3_REQUEST_TYPE_UNKNOWN;
+    aws_s3_request_metrics_get_request_type(metrics, &out_request_type);
+    uint32_t part_number = 0;
+    aws_s3_request_metrics_get_part_number(metrics, &part_number);
+    ASSERT_TRUE(part_number >= 1); /* The part number starts at 1 */
+    /* calculate the expected range start and end based on the part number and part size. */
+    size_t expected_range_start = (part_number - 1) * expected_part_size;
+    size_t expected_range_end = part_number * expected_part_size - 1;
+
+    uint64_t range_start = 0;
+    uint64_t range_end = 0;
+    aws_s3_request_metrics_get_part_range_start(metrics, &range_start);
+    aws_s3_request_metrics_get_part_range_end(metrics, &range_end);
+    ASSERT_UINT_EQUALS(AWS_S3_REQUEST_TYPE_UPLOAD_PART, out_request_type);
+    ASSERT_UINT_EQUALS(expected_range_start, range_start);
+    ASSERT_UINT_EQUALS(expected_range_end, range_end);
+
+    return AWS_OP_SUCCESS;
+}
+
 static int s_check_metrics_helper(
     struct aws_s3_meta_request_test_results *test_results,
     size_t index,
@@ -4996,7 +5023,6 @@ static int s_check_metrics_helper(
     size_t expected_range_end) {
     struct aws_s3_request_metrics *metrics = NULL;
 
-    /* First metrics should be the CreateMPU */
     ASSERT_SUCCESS(aws_array_list_get_at(&test_results->synced_data.succeed_metrics, (void **)&metrics, index));
     enum aws_s3_request_type out_request_type = AWS_S3_REQUEST_TYPE_UNKNOWN;
     aws_s3_request_metrics_get_request_type(metrics, &out_request_type);
@@ -5066,6 +5092,9 @@ static int s_test_s3_round_trip_dynamic_range_size_download_multipart(struct aws
         };
 
         ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &put_options, &test_results));
+        ASSERT_SUCCESS(s_check_metrics_helper(&test_results, 0, AWS_S3_REQUEST_TYPE_CREATE_MULTIPART_UPLOAD, 0, 0, 0));
+        ASSERT_SUCCESS(s_check_metrics_upload_part_helper(&test_results, stored_part_size, 1));
+        ASSERT_SUCCESS(s_check_metrics_upload_part_helper(&test_results, stored_part_size, 2));
         aws_s3_meta_request_test_results_clean_up(&test_results);
 
         /*** GET FILE WITH CHECKSUM -- Head object first ***/