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,11 @@ struct aws_s3_request *aws_s3_request_new(
     uint32_t part_number,
     uint32_t flags);
 
+/* Gets the size of the part payload.
+ * Range is inclusive (i.e. 0-0 range has 1 byte size). In case of misconfigured range returns 0. */
+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_get.c

@@ -253,6 +253,10 @@ static bool s_s3_auto_ranged_get_update(
                             1 /*part_number*/,
                             AWS_S3_REQUEST_FLAG_RECORD_RESPONSE_HEADERS |
                                 AWS_S3_REQUEST_FLAG_ALLOCATE_BUFFER_FROM_POOL);
+                        /* Note: our current default logic is to do part 1, discover size and then abort if payload its
+                         * too huge We optimistically reserve part size for it */
+                        request->part_range_start = 0;
+                        request->part_range_end = meta_request->part_size - 1;
                         ++auto_ranged_get->synced_data.num_parts_requested;
 
                         break;
@@ -285,7 +289,8 @@ static bool s_s3_auto_ranged_get_update(
                         AWS_LOGF_INFO(
                             AWS_LS_S3_META_REQUEST,
                             "id=%p: Doing a ranged get to discover the size of the object and get the first part",
-                            (void *)meta_request);
+                            (void *)meta_request
+                        );
 
                         request = aws_s3_request_new(
                             meta_request,

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"
@@ -1896,7 +1897,10 @@ void s_acquire_mem_and_prepare_request(
     aws_s3_meta_request_prepare_request_callback_fn *callback,
     void *user_data) {
 
-    if (request->ticket == NULL && request->should_allocate_buffer_from_pool) {
+    size_t request_size = (size_t)aws_s3_request_get_part_size(request);
+    AWS_ASSERT(request_size != 0); /* Note: 0 request size is invalid in all cases. */
+
+    if (request->ticket == NULL && request->should_allocate_buffer_from_pool && request_size > 0) {
 
         if (request->send_data.metrics) {
             struct aws_s3_request_metrics *metric = request->send_data.metrics;
@@ -1905,10 +1909,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 = request_size,
         };
 
         struct aws_s3_reserve_memory_payload *payload =

source/s3_meta_request.c

@@ -1429,7 +1429,7 @@ static int s_s3_meta_request_incoming_body(
 
     if (request->send_data.response_body.capacity == 0) {
         /* Make sure that request is get, since puts can also have ticket allocated, which is used for request body. */
-        if (request->request_type == AWS_S3_REQUEST_TYPE_GET_OBJECT && request->ticket != NULL) {
+        if (successful_response && request->request_type == AWS_S3_REQUEST_TYPE_GET_OBJECT && request->ticket != NULL) {
             request->send_data.response_body = aws_s3_buffer_ticket_claim(request->ticket);
         } else {
             size_t buffer_size = s_dynamic_body_initial_buf_size;

source/s3_request.c

@@ -45,6 +45,14 @@ struct aws_s3_request *aws_s3_request_new(
     return request;
 }
 
+uint64_t aws_s3_request_get_part_size(struct aws_s3_request *request) {
+    uint64_t result = 0;
+    if (aws_sub_u64_checked(request->part_range_end, request->part_range_start, &result) != AWS_OP_SUCCESS) {
+        return 0;
+    }
+    return result + 1;
+}
+
 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);