variable first part size

awslabs · Dec 6, 2023 · c0a07fd · c0a07fd
1 parent c7a6e77
commit c0a07fd
Show file tree

Hide file tree

Showing 6 changed files with 72 additions and 47 deletions.
diff --git a/include/aws/s3/private/s3_util.h b/include/aws/s3/private/s3_util.h
@@ -255,8 +255,9 @@ int aws_s3_parse_content_length_response_header(
  * part-ranges on part_size. (ie: if object_range_start is not evenly divisible by part_size, it is considered in the
  * middle of a contiguous part, and that first part will be smaller than part_size.) */
 AWS_S3_API
-uint32_t aws_s3_calculate_auto_range_get_num_parts(
+uint32_t aws_s3_calculate_auto_ranged_get_num_parts(
     size_t part_size,
+    uint64_t first_part_size,
     uint64_t object_range_start,
     uint64_t object_range_end);
 
@@ -276,13 +277,14 @@ int aws_s3_calculate_optimal_mpu_part_size_and_num_parts(
 
 /* Calculates the part range for a part given overall object range, size of each part, and the part's number. Note: part
  * numbers begin at one. This takes into account aligning part-ranges on part_size. Intended to be used in conjunction
- * with aws_s3_calculate_auto_range_get_num_parts. part_number should be less than or equal to the result of
- * aws_s3_calculate_auto_range_get_num_parts. */
+ * with aws_s3_calculate_auto_ranged_get_num_parts. part_number should be less than or equal to the result of
+ * aws_s3_calculate_auto_ranged_get_num_parts. */
 AWS_S3_API
-void aws_s3_calculate_auto_range_get_part_range(
+void aws_s3_calculate_auto_ranged_get_part_range(
     uint64_t object_range_start,
     uint64_t object_range_end,
     size_t part_size,
+    uint64_t first_part_size,
     uint32_t part_number,
     uint64_t *out_part_range_start,
     uint64_t *out_part_range_end);

diff --git a/source/s3_auto_ranged_get.c b/source/s3_auto_ranged_get.c
@@ -155,6 +155,7 @@ static bool s_s3_auto_ranged_get_update(
                  * things simple, we are currently relying on the service to handle turning the Range header into a
                  * Content-Range response header.*/
                 if (auto_ranged_get->initial_message_has_range_header != 0) {
+                    // TODO: align the first part
                     request = aws_s3_request_new(
                         meta_request,
                         AWS_S3_AUTO_RANGE_GET_REQUEST_TYPE_HEAD_OBJECT,
@@ -277,10 +278,11 @@ static bool s_s3_auto_ranged_get_update(
 
                 request->ticket = ticket;
 
-                aws_s3_calculate_auto_range_get_part_range(
+                aws_s3_calculate_auto_ranged_get_part_range(
                     auto_ranged_get->synced_data.object_range_start,
                     auto_ranged_get->synced_data.object_range_end,
                     meta_request->part_size,
+                    auto_ranged_get->synced_data.first_part_size,
                     request->part_number,
                     &request->part_range_start,
                     &request->part_range_end);
@@ -493,16 +495,19 @@ static int s_discover_object_range_and_content_length(
     int error_code,
     uint64_t *out_total_content_length,
     uint64_t *out_object_range_start,
-    uint64_t *out_object_range_end) {
+    uint64_t *out_object_range_end,
+    uint64_t *out_first_part_size) {
     AWS_PRECONDITION(out_total_content_length);
     AWS_PRECONDITION(out_object_range_start);
     AWS_PRECONDITION(out_object_range_end);
+    AWS_PRECONDITION(out_first_part_size);
 
     int result = AWS_OP_ERR;
 
     uint64_t total_content_length = 0;
     uint64_t object_range_start = 0;
     uint64_t object_range_end = 0;
+    uint64_t first_part_size = meta_request->part_size;
 
     AWS_ASSERT(request->discovers_object_size);
     struct aws_s3_auto_ranged_get *auto_ranged_get = meta_request->impl;
@@ -551,10 +556,9 @@ static int s_discover_object_range_and_content_length(
         case AWS_S3_AUTO_RANGE_GET_REQUEST_TYPE_GET_PART_NUMBER:
             AWS_ASSERT(request->part_number == 1);
             AWS_ASSERT(request->send_data.response_headers != NULL);
-            uint64_t first_part_length = 0;
             /* There should be a Content-Length header that indicates the size of first part.*/
             if (aws_s3_parse_content_length_response_header(
-                    meta_request->allocator, request->send_data.response_headers, &first_part_length)) {
+                    meta_request->allocator, request->send_data.response_headers, &first_part_size)) {
 
                 AWS_LOGF_ERROR(
                     AWS_LS_S3_META_REQUEST,
@@ -563,7 +567,7 @@ static int s_discover_object_range_and_content_length(
                     (void *)request);
                 break;
             }
-            if (first_part_length > 0) {
+            if (first_part_size > 0) {
                 /* Parse the object size from the part response. */
                 if (aws_s3_parse_content_range_response_header(
                         meta_request->allocator,
@@ -639,6 +643,7 @@ static int s_discover_object_range_and_content_length(
         *out_total_content_length = total_content_length;
         *out_object_range_start = object_range_start;
         *out_object_range_end = object_range_end;
+        *out_first_part_size = first_part_size;
     }
 
     return result;
@@ -658,6 +663,7 @@ static void s_s3_auto_ranged_get_request_finished(
     uint64_t total_content_length = 0ULL;
     uint64_t object_range_start = 0ULL;
     uint64_t object_range_end = 0ULL;
+    uint64_t first_part_size = 0ULL;
 
     bool found_object_size = false;
     bool request_failed = error_code != AWS_ERROR_SUCCESS;
@@ -667,7 +673,13 @@ static void s_s3_auto_ranged_get_request_finished(
 
         /* Try to discover the object-range and content length.*/
         if (s_discover_object_range_and_content_length(
-                meta_request, request, error_code, &total_content_length, &object_range_start, &object_range_end)) {
+                meta_request,
+                request,
+                error_code,
+                &total_content_length,
+                &object_range_start,
+                &object_range_end,
+                &first_part_size)) {
 
             error_code = aws_last_error_or_unknown();
 
@@ -744,8 +756,14 @@ static void s_s3_auto_ranged_get_request_finished(
             auto_ranged_get->synced_data.object_range_empty = (total_content_length == 0);
             auto_ranged_get->synced_data.object_range_start = object_range_start;
             auto_ranged_get->synced_data.object_range_end = object_range_end;
-            auto_ranged_get->synced_data.total_num_parts = aws_s3_calculate_auto_range_get_num_parts(
-                meta_request->part_size, object_range_start, object_range_end);
+            if (!first_part_size_mismatch) {
+                auto_ranged_get->synced_data.first_part_size = first_part_size;
+            }
+            auto_ranged_get->synced_data.total_num_parts = aws_s3_calculate_auto_ranged_get_num_parts(
+                meta_request->part_size,
+                auto_ranged_get->synced_data.first_part_size,
+                object_range_start,
+                object_range_end);
         }
 
         switch (request->request_tag) {

diff --git a/source/s3_meta_request.c b/source/s3_meta_request.c
@@ -1289,13 +1289,17 @@ static int s_s3_meta_request_headers_block_done(
         uint64_t content_length;
         if (!aws_s3_parse_content_length_response_header(
                 request->allocator, request->send_data.response_headers, &content_length)) {
-            if (content_length > meta_request->part_size ||
-                (content_length < meta_request->part_size &&
-                 aws_http_headers_has(
-                     request->send_data.response_headers,
-                     g_mp_parts_count_header_name) /* Object has multiple parts if this header is present */)) {
+
+            if (content_length > meta_request->part_size * 2) {
                 return aws_raise_error(AWS_ERROR_S3_GET_PART_SIZE_MISMATCH);
             }
+            // if (content_length > meta_request->part_size ||
+            //     (content_length < meta_request->part_size &&
+            //      aws_http_headers_has(
+            //          request->send_data.response_headers,
+            //          g_mp_parts_count_header_name) /* Object has multiple parts if this header is present */)) {
+            //     return aws_raise_error(AWS_ERROR_S3_GET_PART_SIZE_MISMATCH);
+            // }
         }
     }
     return AWS_OP_SUCCESS;

diff --git a/source/s3_util.c b/source/s3_util.c
@@ -485,20 +485,13 @@ int aws_s3_parse_content_length_response_header(
     return result;
 }
 
-uint32_t aws_s3_calculate_auto_range_get_num_parts(
+uint32_t aws_s3_calculate_auto_ranged_get_num_parts(
     size_t part_size,
+    uint64_t first_part_size,
     uint64_t object_range_start,
     uint64_t object_range_end) {
     uint32_t num_parts = 1;
 
-    uint64_t first_part_size = part_size;
-    uint64_t first_part_alignment_offset = object_range_start % part_size;
-
-    /* If the first part size isn't aligned on the assumed part boundary, make it smaller so that it is. */
-    if (first_part_alignment_offset > 0) {
-        first_part_size = part_size - first_part_alignment_offset;
-    }
-
     uint64_t second_part_start = object_range_start + first_part_size;
 
     /* If the range has room for a second part, calculate the additional amount of parts. */
@@ -514,10 +507,11 @@ uint32_t aws_s3_calculate_auto_range_get_num_parts(
     return num_parts;
 }
 
-void aws_s3_calculate_auto_range_get_part_range(
+void aws_s3_calculate_auto_ranged_get_part_range(
     uint64_t object_range_start,
     uint64_t object_range_end,
     size_t part_size,
+    uint64_t first_part_size,
     uint32_t part_number,
     uint64_t *out_part_range_start,
     uint64_t *out_part_range_end) {
@@ -529,16 +523,11 @@ void aws_s3_calculate_auto_range_get_part_range(
     const uint32_t part_index = part_number - 1;
 
     /* Part index is assumed to be in a valid range. */
-    AWS_ASSERT(part_index < aws_s3_calculate_auto_range_get_num_parts(part_size, object_range_start, object_range_end));
+    AWS_ASSERT(
+        part_index <
+        aws_s3_calculate_auto_ranged_get_num_parts(part_size, first_part_size, object_range_start, object_range_end));
 
     uint64_t part_size_uint64 = (uint64_t)part_size;
-    uint64_t first_part_size = part_size_uint64;
-    uint64_t first_part_alignment_offset = object_range_start % part_size_uint64;
-
-    /* Shrink the part to a smaller size if need be to align to the assumed part boundary. */
-    if (first_part_alignment_offset > 0) {
-        first_part_size = part_size_uint64 - first_part_alignment_offset;
-    }
 
     if (part_index == 0) {
         /* If this is the first part, then use the first part size. */

diff --git a/tests/s3_data_plane_tests.c b/tests/s3_data_plane_tests.c
@@ -499,10 +499,11 @@ static int s_test_s3_meta_request_body_streaming(struct aws_allocator *allocator
             struct aws_s3_request *request = aws_s3_request_new(
                 meta_request, 0 /*request_tag*/, AWS_S3_REQUEST_TYPE_GET_OBJECT, part_number, 0 /*flags*/);
 
-            aws_s3_calculate_auto_range_get_part_range(
+            aws_s3_calculate_auto_ranged_get_part_range(
                 0ULL,
                 total_object_size - 1,
-                (uint64_t)request_response_body_size,
+                request_response_body_size /*part_size*/,
+                (uint64_t)request_response_body_size /*first_part_size*/,
                 part_number,
                 &request->part_range_start,
                 &request->part_range_end);
@@ -536,10 +537,11 @@ static int s_test_s3_meta_request_body_streaming(struct aws_allocator *allocator
             struct aws_s3_request *request = aws_s3_request_new(
                 meta_request, 0 /*request_tag*/, AWS_S3_REQUEST_TYPE_GET_OBJECT, part_number, 0 /*flags*/);
 
-            aws_s3_calculate_auto_range_get_part_range(
+            aws_s3_calculate_auto_ranged_get_part_range(
                 0ULL,
                 total_object_size - 1,
-                (uint64_t)request_response_body_size,
+                request_response_body_size /*part_size*/,
+                (uint64_t)request_response_body_size /*first_part_size*/,
                 part_number,
                 &request->part_range_start,
                 &request->part_range_end);
@@ -565,10 +567,11 @@ static int s_test_s3_meta_request_body_streaming(struct aws_allocator *allocator
         struct aws_s3_request *request = aws_s3_request_new(
             meta_request, 0 /*request_tag*/, AWS_S3_REQUEST_TYPE_GET_OBJECT, part_range1_start, 0 /*flags*/);
 
-        aws_s3_calculate_auto_range_get_part_range(
+        aws_s3_calculate_auto_ranged_get_part_range(
             0ULL,
             total_object_size - 1,
-            (uint64_t)request_response_body_size,
+            request_response_body_size /*part_size*/,
+            (uint64_t)request_response_body_size /*first_part_size*/,
             part_range1_start,
             &request->part_range_start,
             &request->part_range_end);

diff --git a/tests/s3_util_tests.c b/tests/s3_util_tests.c
@@ -264,8 +264,14 @@ static int s_validate_part_ranges(
         uint64_t part_range_start = 0ULL;
         uint64_t part_range_end = 0ULL;
 
-        aws_s3_calculate_auto_range_get_part_range(
-            object_range_start, object_range_end, part_size, i + 1, &part_range_start, &part_range_end);
+        aws_s3_calculate_auto_ranged_get_part_range(
+            object_range_start,
+            object_range_end,
+            part_size,
+            part_size - (object_range_start % part_size),
+            i + 1,
+            &part_range_start,
+            &part_range_end);
 
         ASSERT_TRUE(part_range_start == part_ranges[i * 2]);
         ASSERT_TRUE(part_range_end == part_ranges[i * 2 + 1]);
@@ -297,7 +303,7 @@ static int s_test_s3_get_num_parts_and_get_part_range(struct aws_allocator *allo
         };
 
         ASSERT_TRUE(
-            aws_s3_calculate_auto_range_get_num_parts(part_size, object_range_start, object_range_end) ==
+            aws_s3_calculate_auto_ranged_get_num_parts(part_size, part_size, object_range_start, object_range_end) ==
             expected_num_parts);
 
         ASSERT_SUCCESS(
@@ -323,7 +329,8 @@ static int s_test_s3_get_num_parts_and_get_part_range(struct aws_allocator *allo
         };
 
         ASSERT_TRUE(
-            aws_s3_calculate_auto_range_get_num_parts(part_size, object_range_start, object_range_end) ==
+            aws_s3_calculate_auto_ranged_get_num_parts(
+                part_size, part_size - (object_range_start % part_size), object_range_start, object_range_end) ==
             expected_num_parts);
 
         ASSERT_SUCCESS(
@@ -353,7 +360,8 @@ static int s_test_s3_get_num_parts_and_get_part_range(struct aws_allocator *allo
         };
 
         ASSERT_TRUE(
-            aws_s3_calculate_auto_range_get_num_parts(part_size, object_range_start, object_range_end) ==
+            aws_s3_calculate_auto_ranged_get_num_parts(
+                part_size, part_size - (object_range_start % part_size), object_range_start, object_range_end) ==
             expected_num_parts);
 
         ASSERT_SUCCESS(
@@ -369,7 +377,8 @@ static int s_test_s3_get_num_parts_and_get_part_range(struct aws_allocator *allo
         const uint64_t part_ranges[] = {8, 8};
 
         ASSERT_TRUE(
-            aws_s3_calculate_auto_range_get_num_parts(part_size, object_range_start, object_range_end) ==
+            aws_s3_calculate_auto_ranged_get_num_parts(
+                part_size, part_size - (object_range_start % part_size), object_range_start, object_range_end) ==
             expected_num_parts);
 
         ASSERT_SUCCESS(