wip async reserve

DmitriyMusatkin · DmitriyMusatkin · commit 6d9cc7961ac3 · 2025-04-23T01:15:20.000-07:00
diff --git a/include/aws/s3/private/s3_default_buffer_pool.h b/include/aws/s3/private/s3_default_buffer_pool.h
@@ -7,6 +7,7 @@
  */
 
 #include <aws/s3/s3.h>
+#include <aws/s3/s3_buffer_pool.h>
 
 /*
  * S3 buffer pool.
@@ -95,19 +96,9 @@ AWS_S3_API void aws_s3_default_buffer_pool_destroy(struct aws_s3_default_buffer_
  * If you MUST acquire a buffer now (waiting to reserve a ticket would risk deadlock),
  * use aws_s3_buffer_pool_acquire_forced_buffer() instead.
  */
-AWS_S3_API struct aws_s3_buffer_ticket *aws_s3_default_buffer_pool_reserve(
+AWS_S3_API struct aws_future_s3_buffer_ticket *aws_s3_default_buffer_pool_reserve(
     struct aws_s3_default_buffer_pool *buffer_pool,
-    size_t size);
-
-/*
- * Whether pool has a reservation hold.
- */
-AWS_S3_API bool aws_s3_default_buffer_pool_has_reservation_hold(struct aws_s3_default_buffer_pool *buffer_pool);
-
-/*
- * Remove reservation hold on pool.
- */
-AWS_S3_API void aws_s3_default_buffer_pool_remove_reservation_hold(struct aws_s3_default_buffer_pool *buffer_pool);
+    struct aws_s3_buffer_pool_reserve_meta meta);
 
 /*
  * Trades in the ticket for a buffer.
diff --git a/include/aws/s3/private/s3_meta_request_impl.h b/include/aws/s3/private/s3_meta_request_impl.h
@@ -46,7 +46,9 @@ typedef void(aws_s3_meta_request_prepare_request_callback_fn)(
 struct aws_s3_prepare_request_payload {
     struct aws_allocator *allocator;
     struct aws_s3_request *request;
+    struct aws_event_loop *event_loop;
     struct aws_task task;
+    struct aws_future_s3_buffer_ticket *async_buffer_reserve;
     /* async step: wait for vtable->prepare_request() call to complete */
     struct aws_future_void *asyncstep_prepare_request;
     /* callback to invoke when all request preparation work is complete */
diff --git a/include/aws/s3/private/s3_request.h b/include/aws/s3/private/s3_request.h
@@ -20,9 +20,8 @@ struct aws_s3_meta_request;
 
 enum aws_s3_request_flags {
     AWS_S3_REQUEST_FLAG_RECORD_RESPONSE_HEADERS = 0x00000001,
-    AWS_S3_REQUEST_FLAG_PART_SIZE_RESPONSE_BODY = 0x00000002,
-    AWS_S3_REQUEST_FLAG_ALWAYS_SEND = 0x00000004,
-    AWS_S3_REQUEST_FLAG_PART_SIZE_REQUEST_BODY = 0x00000008,
+    AWS_S3_REQUEST_FLAG_ALWAYS_SEND = 0x00000002,
+    AWS_S3_REQUEST_FLAG_ALLOCATE_BUFFER_FROM_POOL = 0x00000004,
 };
 
 /**
@@ -224,11 +223,8 @@ struct aws_s3_request {
     /* When true, response headers from the request will be stored in the request's response_headers variable. */
     uint32_t record_response_headers : 1;
 
-    /* When true, the response body buffer will be allocated in the size of a part. */
-    uint32_t has_part_size_response_body : 1;
-
-    /* When true, the request body buffer will be allocated in the size of a part. */
-    uint32_t has_part_size_request_body : 1;
+    /* Indicates whether buffer should be allocated for the request from the pool. */
+    uint32_t should_allocate_buffer_from_pool : 1;
 
     /* When true, this request is being tracked by the client for limiting the amount of in-flight-requests/stats. */
     uint32_t tracked_by_client : 1;
diff --git a/include/aws/s3/s3.h b/include/aws/s3/s3.h
@@ -50,6 +50,7 @@ enum aws_s3_errors {
     AWS_ERROR_S3_RECV_FILE_NOT_FOUND,
     AWS_ERROR_S3_REQUEST_TIMEOUT,
     AWS_ERROR_S3_TOKEN_EXPIRED,
+    AWS_ERROR_S3_BUFFER_ALLOCATION_FAILED,
 
     AWS_ERROR_S3_END_RANGE = AWS_ERROR_ENUM_END_RANGE(AWS_C_S3_PACKAGE_ID)
 };
diff --git a/include/aws/s3/s3_buffer_pool.h b/include/aws/s3/s3_buffer_pool.h
@@ -0,0 +1,48 @@
+#ifndef AWS_S3_BUFFER_POOL_H
+#define AWS_S3_BUFFER_POOL_H
+
+#include <aws/io/future.h>
+#include <aws/s3/s3.h>
+
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+AWS_PUSH_SANE_WARNING_LEVEL
+AWS_EXTERN_C_BEGIN
+struct aws_s3_buffer_ticket;
+
+/**
+ * aws_future<aws_s3_buffer_ticket*>
+ */
+AWS_FUTURE_T_POINTER_WITH_RELEASE_DECLARATION(aws_future_s3_buffer_ticket, struct aws_s3_buffer_ticket, AWS_S3_API)
+
+struct aws_s3_buffer_pool_reserve_meta {
+    struct aws_s3_client *client;
+    struct aws_s3_meta_request *meta_request;
+    size_t size;
+    bool can_block;
+};
+
+struct aws_s3_buffer_pool {
+    void (*destroy)(struct aws_s3_buffer_pool *pool);
+
+    struct aws_future_s3_buffer_ticket *(*reserve)(struct aws_s3_buffer_pool *pool, 
+        struct aws_s3_buffer_pool_reserve_meta meta);
+    struct aws_byte_buf (*claim)(struct aws_s3_buffer_pool *pool, struct aws_s3_buffer_ticket *ticket);
+    void (*release)(struct aws_s3_buffer_pool *pool, struct aws_s3_buffer_ticket *ticket);
+
+    void (*trim)(struct aws_s3_buffer_pool *pool);
+
+    struct aws_allocator *allocator;
+    void *user_data;
+};
+
+typedef struct aws_s3_buffer_pool *(aws_s3_buffer_pool_factory_fn)(struct aws_allocator *allocator,
+                                                                   uint64_t part_size,
+                                                                   uint64_t mem_limit);
+
+AWS_EXTERN_C_END
+AWS_POP_SANE_WARNING_LEVEL
+ 
+ #endif /* AWS_S3_BUFFER_POOL_H */
diff --git a/include/aws/s3/s3_client.h b/include/aws/s3/s3_client.h
@@ -10,6 +10,7 @@
 #include <aws/common/ref_count.h>
 #include <aws/io/retry_strategy.h>
 #include <aws/s3/s3.h>
+#include <aws/s3/s3_buffer_pool.h>
 
 AWS_PUSH_SANE_WARNING_LEVEL
 
@@ -373,25 +374,6 @@ typedef struct aws_s3express_credentials_provider *(
                                        void *shutdown_user_data,
                                        void *factory_user_data);
 
-struct aws_s3_buffer_ticket;
-
-struct aws_s3_buffer_pool {
-    void (*destroy)(struct aws_s3_buffer_pool *pool);
-
-    struct aws_s3_buffer_ticket *(*reserve)(struct aws_s3_buffer_pool *pool, size_t size);
-    struct aws_byte_buf (*claim)(struct aws_s3_buffer_pool *pool, struct aws_s3_buffer_ticket *ticket);
-    void (*release)(struct aws_s3_buffer_pool *pool, struct aws_s3_buffer_ticket *ticket);
-
-    void (*trim)(struct aws_s3_buffer_pool *pool);
-
-    struct aws_allocator *allocator;
-    void *user_data;
-};
-
-typedef struct aws_s3_buffer_pool *(aws_s3_buffer_pool_factory_fn)(struct aws_allocator *allocator,
-                                                                   uint64_t part_size,
-                                                                   uint64_t mem_limit);
-
 /* Keepalive properties are TCP only.
  * If interval or timeout are zero, then default values are used.
  */
diff --git a/source/s3.c b/source/s3.c
@@ -53,6 +53,7 @@ static struct aws_error_info s_errors[] = {
     AWS_DEFINE_ERROR_INFO_S3(AWS_ERROR_S3_RECV_FILE_NOT_FOUND, "The receive file doesn't exist, cannot create as configuration required."),
     AWS_DEFINE_ERROR_INFO_S3(AWS_ERROR_S3_REQUEST_TIMEOUT, "RequestTimeout error received from S3."),
     AWS_DEFINE_ERROR_INFO_S3(AWS_ERROR_S3_TOKEN_EXPIRED, "Token expired - needs a refresh."),
+    AWS_DEFINE_ERROR_INFO_S3(AWS_ERROR_S3_BUFFER_ALLOCATION_FAILED, "Could not reserve memory for object."),
 };
 /* clang-format on */
 
diff --git a/source/s3_auto_ranged_get.c b/source/s3_auto_ranged_get.c
@@ -225,7 +225,6 @@ static bool s_s3_auto_ranged_get_update(
                     auto_ranged_get->synced_data.num_parts_requested > 0) {
                     goto has_work_remaining;
                 }
-                struct aws_s3_buffer_ticket *ticket = NULL;
                 switch (s_s3_get_request_type_for_discovering_object_size(meta_request)) {
                     case AWS_S3_AUTO_RANGE_GET_REQUEST_TYPE_HEAD_OBJECT:
                         AWS_LOGF_INFO(
@@ -246,20 +245,13 @@ static bool s_s3_auto_ranged_get_update(
                             "id=%p: Doing a 'GET_OBJECT_WITH_PART_NUMBER_1' to discover the size of the object and get "
                             "the first part",
                             (void *)meta_request);
-                        ticket = meta_request->client->buffer_pool->reserve(
-                            meta_request->client->buffer_pool, meta_request->part_size);
-
-                        if (ticket == NULL) {
-                            goto has_work_remaining;
-                        }
 
                         request = aws_s3_request_new(
                             meta_request,
                             AWS_S3_AUTO_RANGE_GET_REQUEST_TYPE_GET_OBJECT_WITH_PART_NUMBER_1,
                             AWS_S3_REQUEST_TYPE_GET_OBJECT,
                             1 /*part_number*/,
-                            AWS_S3_REQUEST_FLAG_RECORD_RESPONSE_HEADERS | AWS_S3_REQUEST_FLAG_PART_SIZE_RESPONSE_BODY);
-                        request->ticket = ticket;
+                            AWS_S3_REQUEST_FLAG_RECORD_RESPONSE_HEADERS | AWS_S3_REQUEST_FLAG_ALLOCATE_BUFFER_FROM_POOL);
                         ++auto_ranged_get->synced_data.num_parts_requested;
 
                         break;
@@ -294,28 +286,12 @@ static bool s_s3_auto_ranged_get_update(
                             "id=%p: Doing a ranged get to discover the size of the object and get the first part",
                             (void *)meta_request);
 
-                        if (first_part_size >= s_min_size_response_for_pooling) {
-                            /* Note: explicitly reserving the whole part size
-                             * even if expect to receive less data. Pool will
-                             * reserve the whole part size for it anyways, so no
-                             * reason getting a smaller chunk. */
-                            ticket = meta_request->client->buffer_pool->reserve(
-                                meta_request->client->buffer_pool, (size_t)meta_request->part_size);
-
-                            if (ticket == NULL) {
-                                goto has_work_remaining;
-                            }
-                        } else {
-                            ticket = NULL;
-                        }
-
                         request = aws_s3_request_new(
                             meta_request,
                             AWS_S3_AUTO_RANGE_GET_REQUEST_TYPE_GET_OBJECT_WITH_RANGE,
                             AWS_S3_REQUEST_TYPE_GET_OBJECT,
                             1 /*part_number*/,
-                            AWS_S3_REQUEST_FLAG_RECORD_RESPONSE_HEADERS | AWS_S3_REQUEST_FLAG_PART_SIZE_RESPONSE_BODY);
-                        request->ticket = ticket;
+                            AWS_S3_REQUEST_FLAG_RECORD_RESPONSE_HEADERS | AWS_S3_REQUEST_FLAG_ALLOCATE_BUFFER_FROM_POOL);
                         request->part_range_start = part_range_start;
                         request->part_range_end = part_range_start + first_part_size - 1; /* range-end is inclusive */
                         ++auto_ranged_get->synced_data.num_parts_requested;
@@ -361,21 +337,12 @@ static bool s_s3_auto_ranged_get_update(
                     auto_ranged_get->synced_data.read_window_warning_issued = 0;
                 }
 
-                struct aws_s3_buffer_ticket *ticket = meta_request->client->buffer_pool->reserve(
-                    meta_request->client->buffer_pool, meta_request->part_size);
-
-                if (ticket == NULL) {
-                    goto has_work_remaining;
-                }
-
                 request = aws_s3_request_new(
                     meta_request,
                     AWS_S3_AUTO_RANGE_GET_REQUEST_TYPE_GET_OBJECT_WITH_RANGE,
                     AWS_S3_REQUEST_TYPE_GET_OBJECT,
                     auto_ranged_get->synced_data.num_parts_requested + 1 /*part_number*/,
-                    AWS_S3_REQUEST_FLAG_PART_SIZE_RESPONSE_BODY);
-
-                request->ticket = ticket;
+                    AWS_S3_REQUEST_FLAG_ALLOCATE_BUFFER_FROM_POOL);
 
                 aws_s3_calculate_auto_ranged_get_part_range(
                     auto_ranged_get->synced_data.object_range_start,
@@ -530,11 +497,11 @@ static struct aws_future_void *s_s3_auto_ranged_get_prepare_request(struct aws_s
     /* Success! */
     AWS_LOGF_DEBUG(
         AWS_LS_S3_META_REQUEST,
-        "id=%p: Created request %p for part %d part sized %d",
+        "id=%p: Created request %p for part %d allocated from pool %d",
         (void *)meta_request,
         (void *)request,
         request->part_number,
-        request->has_part_size_response_body);
+        request->should_allocate_buffer_from_pool);
 
     success = true;
 
diff --git a/source/s3_auto_ranged_put.c b/source/s3_auto_ranged_put.c
@@ -567,20 +567,21 @@ static bool s_s3_auto_ranged_put_update(
                     AWS_FATAL_ASSERT(meta_request->synced_data.async_write.buffered_data_ticket);
                     ticket = meta_request->synced_data.async_write.buffered_data_ticket;
                     meta_request->synced_data.async_write.buffered_data_ticket = NULL;
-                } else {
-                    /* Try to reserve a ticket */
-                    ticket = meta_request->client->buffer_pool->reserve(
-                        meta_request->client->buffer_pool, meta_request->part_size);
                 }
 
                 if (ticket != NULL) {
                     /* Allocate a request for another part. */
+                    uint32_t flags = AWS_S3_REQUEST_FLAG_RECORD_RESPONSE_HEADERS;
+                    if (meta_request->synced_data.async_write.ready_to_send) {
+                        flags |= AWS_S3_REQUEST_FLAG_ALLOCATE_BUFFER_FROM_POOL;
+                    }
+
                     request = aws_s3_request_new(
                         meta_request,
                         AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_PART,
                         AWS_S3_REQUEST_TYPE_UPLOAD_PART,
                         0 /*part_number*/,
-                        AWS_S3_REQUEST_FLAG_RECORD_RESPONSE_HEADERS | AWS_S3_REQUEST_FLAG_PART_SIZE_REQUEST_BODY);
+                        flags);
 
                     request->part_number = auto_ranged_put->threaded_update_data.next_part_number;
 
diff --git a/source/s3_client.c b/source/s3_client.c
@@ -232,10 +232,11 @@ static void s_default_buffer_pool_destroy(struct aws_s3_buffer_pool *pool) {
     aws_mem_release(pool->allocator, pool);
 }
 
-struct aws_s3_buffer_ticket *s_default_pool_reserve(struct aws_s3_buffer_pool *pool, size_t size) {
+struct aws_future_s3_buffer_ticket *s_default_pool_reserve(struct aws_s3_buffer_pool *pool, 
+        struct aws_s3_buffer_pool_reserve_meta meta) {
     struct aws_s3_default_buffer_pool *default_pool = (struct aws_s3_default_buffer_pool *)pool->user_data;
 
-    return aws_s3_default_buffer_pool_reserve(default_pool, size);
+    return aws_s3_default_buffer_pool_reserve(default_pool, meta);
 }
 
 struct aws_byte_buf s_default_pool_claim(struct aws_s3_buffer_pool *pool, struct aws_s3_buffer_ticket *ticket) {
@@ -1868,11 +1869,6 @@ void aws_s3_client_update_meta_requests_threaded(struct aws_s3_client *client) {
 
     const uint32_t num_passes = AWS_ARRAY_SIZE(pass_flags);
 
-    /*
-    Do we still need it?
-    aws_s3_buffer_pool_remove_reservation_hold(client->buffer_pool);
-    */
-
     for (uint32_t pass_index = 0; pass_index < num_passes; ++pass_index) {
 
         /**
diff --git a/source/s3_default_buffer_pool.c b/source/s3_default_buffer_pool.c
diff --git a/source/s3_meta_request.c b/source/s3_meta_request.c
diff --git a/source/s3_request.c b/source/s3_request.c
diff --git a/source/s3_request_messages.c b/source/s3_request_messages.c
