Optimize the sizes of buffers requested from mem pool (#563)

Author: DmitriyMusatkin

.github/workflows/ci.yml

@@ -37,7 +37,7 @@ jobs:
           - manylinux2014-x86
           - al2-x64
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -88,7 +88,7 @@ jobs:
           # our base Linux image to Ubuntu 22.
             extra-build-flag: --cmake-extra=-DCMAKE_EXE_LINKER_FLAGS="-lresolv"
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -104,7 +104,7 @@ jobs:
       matrix:
         sanitizers: [",thread", ",address,undefined"]
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -117,7 +117,7 @@ jobs:
   linux-shared-libs:
     runs-on: ubuntu-24.04 # latest
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -130,7 +130,7 @@ jobs:
   byo-crypto:
     runs-on: ubuntu-24.04 # latest
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -143,7 +143,7 @@ jobs:
   windows:
     runs-on: windows-2025 # latest
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -159,7 +159,7 @@ jobs:
       matrix:
         arch: [x86, x64]
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -171,7 +171,7 @@ jobs:
   windows-shared-libs:
     runs-on: windows-2025 # latest
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -183,7 +183,7 @@ jobs:
   windows-app-verifier:
     runs-on: windows-2025 # latest
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -198,7 +198,7 @@ jobs:
   macos:
     runs-on: macos-14 # latest
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -215,7 +215,7 @@ jobs:
   macos-x64:
     runs-on: macos-14-large # latest
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -234,7 +234,7 @@ jobs:
   downstream:
     runs-on: ubuntu-24.04 # latest
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}
@@ -247,7 +247,7 @@ jobs:
   linux-debug:
     runs-on: ubuntu-24.04 # latest
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}

.github/workflows/codecov.yml

@@ -20,7 +20,7 @@ jobs:
   codecov-linux:
     runs-on: ubuntu-24.04
     steps:
-    - uses: aws-actions/configure-aws-credentials@v4
+    - uses: aws-actions/configure-aws-credentials@v6
       with:
         role-to-assume: ${{ env.CRT_CI_ROLE }}
         aws-region: ${{ env.AWS_DEFAULT_REGION }}

include/aws/s3/private/s3_default_buffer_pool.h

@@ -44,6 +44,9 @@ struct aws_s3_default_buffer_pool_usage_stats {
     /* Max size of buffer to be allocated from primary. */
     size_t primary_cutoff;
 
+    /* Min size of buffer to be allocated from primary. */
+    size_t primary_min_cutoff;
+
     /* Overall memory allocated for blocks. */
     size_t primary_allocated;
     /* Number of blocks allocated in primary. */
@@ -91,6 +94,9 @@ struct aws_s3_default_buffer_pool {
     /* size at which allocations should go to secondary */
     size_t primary_size_cutoff;
 
+    /* size below which allocations should go to secondary */
+    size_t primary_size_min_cutoff;
+
     /* NOTE: See aws_s3_buffer_pool_usage_stats for descriptions of most fields */
 
     size_t mem_limit;
@@ -185,6 +191,18 @@ AWS_S3_API struct aws_byte_buf aws_s3_default_buffer_pool_acquire_buffer(
 AWS_S3_API struct aws_s3_default_buffer_pool_usage_stats aws_s3_default_buffer_pool_get_usage(
     struct aws_s3_buffer_pool *buffer_pool);
 
+enum aws_s3_default_buffer_pool_reserved_from {
+    AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY = 0,
+    AWS_S3_BUFFER_POOL_RESERVED_FROM_PRIMARY,
+    AWS_S3_BUFFER_POOL_RESERVED_FROM_SPECIAL,
+};
+
+/*
+ * Helper to determine which area of default buffer pool ticket was reserved from.
+ */
+AWS_S3_API enum aws_s3_default_buffer_pool_reserved_from aws_s3_default_buffer_pool_get_ticket_reserved_from(
+    struct aws_s3_default_buffer_ticket *ticket);
+
 /*
  * Trims all unused mem from the pool.
  * Warning: fairly slow operation, do not use in critical path.

include/aws/s3/private/s3_request.h

@@ -240,10 +240,12 @@ struct aws_s3_request {
      */
     struct aws_s3_buffer_ticket *ticket;
 
-    /* Beginning range of this part. */
+    /* Beginning range of this part.  */
     uint64_t part_range_start;
 
-    /* Last byte of this part.*/
+    /* 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.)
@@ -364,6 +366,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_payload_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

@@ -265,6 +265,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;

source/s3_auto_ranged_put.c

@@ -323,6 +323,36 @@ static struct aws_s3_meta_request_vtable s_s3_auto_ranged_put_vtable = {
     .pause = s_s3_auto_ranged_put_pause,
 };
 
+/**
+ * Helper to initialize the request ranges and content-length
+ * based on the request->part_number and meta_request->part_size
+ */
+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 &&
+        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;
+        }
+    }
+    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;
+}
+
 /* 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,
@@ -605,6 +635,8 @@ static bool s_s3_auto_ranged_put_update(
 
                 request->part_number = auto_ranged_put->threaded_update_data.next_part_number;
 
+                s_compute_request_body_size(meta_request, request);
+
                 /* 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;
@@ -761,36 +793,6 @@ static bool s_s3_auto_ranged_put_update(
     return work_remaining;
 }
 
-/**
- * Helper to initialize the request ranges and content-length
- * based on the request->part_number and meta_request->part_size
- */
-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 &&
-        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;
-        }
-    }
-    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(
     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"
@@ -2059,7 +2060,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_payload_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;
@@ -2069,10 +2073,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 = request->buffer_size,
+            .size = request_size,
         };
 
         struct aws_s3_reserve_memory_payload *payload =