Optimize the sizes of buffers requested from mem pool

include/aws/s3/private/s3_request.h

@@ -166,12 +166,12 @@ 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. */
+    /* Beginning range of this part.  */
     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. */
+    /* Last byte of this part.
+     * Note: this is available on both put and get, but in put case it can be optimistic,
+     * ex for unknown content length we dont know what the end will be until we finish reading data. */
     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.)
@@ -292,6 +292,12 @@ struct aws_s3_request *aws_s3_request_new(
     uint32_t part_number,
     uint32_t flags);
 
+/* Gets the size of the part payload.
+ * returns metarequest part size if accurate size is not available for request that have part size body or response or 0
+ * otherwise. */
+AWS_S3_API
+uint64_t aws_s3_request_get_part_size(struct aws_s3_request *request);
+
 /* Set up the request to be sent. Called each time before the request is sent. Will initially call
  * aws_s3_request_clean_up_send_data to clear out anything previously existing in send_data. */
 AWS_S3_API

source/s3_auto_ranged_put.c

@@ -322,6 +322,34 @@ static struct aws_s3_meta_request_vtable s_s3_auto_ranged_put_vtable = {
     .pause = s_s3_auto_ranged_put_pause,
 };
 
+/**
+ * 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.
+ */
+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);
+
+    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) {
+        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;
+}
+
 /* Allocate a new auto-ranged put meta request */
 struct aws_s3_meta_request *aws_s3_meta_request_auto_ranged_put_new(
     struct aws_allocator *allocator,
@@ -592,6 +620,10 @@ static bool s_s3_auto_ranged_put_update(
 
                 request->part_number = auto_ranged_put->threaded_update_data.next_part_number;
 
+                size_t request_body_size =
+                    s_compute_request_body_size(meta_request, request->part_number, &request->part_range_start);
+                request->part_range_end = request->part_range_start + request_body_size - 1;
+
                 /* If request was previously uploaded, we prepare it to ensure checksums still match,
                  * but ultimately it gets marked no-op and we don't send it */
                 request->was_previously_uploaded = request_previously_uploaded;
@@ -748,34 +780,6 @@ static bool s_s3_auto_ranged_put_update(
     return work_remaining;
 }
 
-/**
- * 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.
- */
-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);
-
-    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) {
-        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;
-}
-
 static int s_verify_part_matches_checksum(
     struct aws_allocator *allocator,
     struct aws_byte_cursor body_cur,

source/s3_client.c

@@ -11,6 +11,7 @@
 #include "aws/s3/private/s3_default_meta_request.h"
 #include "aws/s3/private/s3_meta_request_impl.h"
 #include "aws/s3/private/s3_parallel_input_stream.h"
+#include "aws/s3/private/s3_request.h"
 #include "aws/s3/private/s3_request_messages.h"
 #include "aws/s3/private/s3_util.h"
 #include "aws/s3/private/s3express_credentials_provider_impl.h"
@@ -1905,10 +1906,11 @@ void s_acquire_mem_and_prepare_request(
 
         struct aws_allocator *allocator = request->allocator;
         struct aws_s3_meta_request *meta_request = request->meta_request;
+
         struct aws_s3_buffer_pool_reserve_meta meta = {
             .client = client,
             .meta_request = meta_request,
-            .size = meta_request->part_size,
+            .size = (size_t)aws_s3_request_get_part_size(request),
         };
 
         struct aws_s3_reserve_memory_payload *payload =

source/s3_request.c

@@ -45,6 +45,25 @@ struct aws_s3_request *aws_s3_request_new(
     return request;
 }
 
+uint64_t aws_s3_request_get_part_size(struct aws_s3_request *request) {
+    /* Note: this is still somewhat suboptimal as 0-0 is a valid 1 byte range, but
+     * we will overallocate for it */
+    bool is_range_set = request->part_range_end != request->part_range_start || request->part_range_start != 0;
+
+    if (is_range_set) {
+        /* +1 cause range end is inclusive */
+        return request->part_range_end - request->part_range_start + 1;
+    }
+
+    if (request->request_type == AWS_S3_REQUEST_TYPE_GET_OBJECT ||
+        request->request_type == AWS_S3_REQUEST_TYPE_PUT_OBJECT ||
+        request->request_type == AWS_S3_REQUEST_TYPE_UPLOAD_PART) {
+        return request->meta_request->part_size;
+    }
+
+    return 0;
+}
+
 static void s_populate_metrics_from_message(struct aws_s3_request *request, struct aws_http_message *message) {
     struct aws_byte_cursor out_path;
     AWS_ZERO_STRUCT(out_path);