Memory pool interface (#517)

Co-authored-by: Waqar Ahmed Khan <[email protected]>

Author: DmitriyMusatkin

include/aws/s3/private/s3_buffer_pool.h …/aws/s3/private/s3_default_buffer_pool.hinclude/aws/s3/private/s3_buffer_pool.h renamed to include/aws/s3/private/s3_default_buffer_pool.h include/aws/s3/private/s3_buffer_pool.h renamed to 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.
@@ -30,10 +31,10 @@
 
 AWS_EXTERN_C_BEGIN
 
-struct aws_s3_buffer_pool;
-struct aws_s3_buffer_pool_ticket;
+struct aws_s3_default_buffer_pool;
+struct aws_s3_default_buffer_ticket;
 
-struct aws_s3_buffer_pool_usage_stats {
+struct aws_s3_default_buffer_pool_usage_stats {
     /* Effective Max memory limit. Memory limit value provided during construction minus
      * buffer reserved for overhead of the pool */
     size_t mem_limit;
@@ -71,16 +72,15 @@ struct aws_s3_buffer_pool_usage_stats {
  * buffers can no longer be reserved from (reservation hold is placed on the pool).
  * Returns buffer pool pointer on success and NULL on failure.
  */
-AWS_S3_API struct aws_s3_buffer_pool *aws_s3_buffer_pool_new(
+AWS_S3_API struct aws_s3_buffer_pool *aws_s3_default_buffer_pool_new(
     struct aws_allocator *allocator,
-    size_t chunk_size,
-    size_t mem_limit);
+    struct aws_s3_buffer_pool_config config);
 
 /*
  * Destroys buffer pool.
  * Does nothing if buffer_pool is NULL.
  */
-AWS_S3_API void aws_s3_buffer_pool_destroy(struct aws_s3_buffer_pool *buffer_pool);
+AWS_S3_API void aws_s3_default_buffer_pool_destroy(struct aws_s3_buffer_pool *buffer_pool);
 
 /*
  * Reserves memory from the pool for later use.
@@ -95,60 +95,32 @@ AWS_S3_API void aws_s3_buffer_pool_destroy(struct aws_s3_buffer_pool *buffer_poo
  * 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_pool_ticket *aws_s3_buffer_pool_reserve(
+AWS_S3_API struct aws_future_s3_buffer_ticket *aws_s3_default_buffer_pool_reserve(
     struct aws_s3_buffer_pool *buffer_pool,
-    size_t size);
-
-/*
- * Whether pool has a reservation hold.
- */
-AWS_S3_API bool aws_s3_buffer_pool_has_reservation_hold(struct aws_s3_buffer_pool *buffer_pool);
-
-/*
- * Remove reservation hold on pool.
- */
-AWS_S3_API void aws_s3_buffer_pool_remove_reservation_hold(struct aws_s3_buffer_pool *buffer_pool);
+    struct aws_s3_buffer_pool_reserve_meta meta);
 
 /*
  * Trades in the ticket for a buffer.
  * Cannot fail and can over allocate above mem limit if reservation was not accurate.
  * Using the same ticket twice will return the same buffer.
  * Buffer is only valid until the ticket is released.
  */
-AWS_S3_API struct aws_byte_buf aws_s3_buffer_pool_acquire_buffer(
-    struct aws_s3_buffer_pool *buffer_pool,
-    struct aws_s3_buffer_pool_ticket *ticket);
-
-/*
- * Force immediate acquisition of a buffer from the pool.
- * This should only be used if waiting to reserve a ticket would risk deadlock.
- * This cannot fail, not even if the pool has a reservation hold,
- * not even if the memory limit has been exceeded.
- */
-AWS_S3_API struct aws_byte_buf aws_s3_buffer_pool_acquire_forced_buffer(
-    struct aws_s3_buffer_pool *buffer_pool,
-    size_t size,
-    struct aws_s3_buffer_pool_ticket **out_new_ticket);
-
-/*
- * Releases the ticket.
- * Any buffers associated with the ticket are invalidated.
- */
-AWS_S3_API void aws_s3_buffer_pool_release_ticket(
+AWS_S3_API struct aws_byte_buf aws_s3_default_buffer_pool_acquire_buffer(
     struct aws_s3_buffer_pool *buffer_pool,
-    struct aws_s3_buffer_pool_ticket *ticket);
+    struct aws_s3_default_buffer_ticket *ticket);
 
 /*
  * Get pool memory usage stats.
  */
-AWS_S3_API struct aws_s3_buffer_pool_usage_stats aws_s3_buffer_pool_get_usage(struct aws_s3_buffer_pool *buffer_pool);
+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);
 
 /*
  * Trims all unused mem from the pool.
  * Warning: fairly slow operation, do not use in critical path.
  * TODO: partial trimming? ex. only trim down to 50% of max?
  */
-AWS_S3_API void aws_s3_buffer_pool_trim(struct aws_s3_buffer_pool *buffer_pool);
+AWS_S3_API void aws_s3_default_buffer_pool_trim(struct aws_s3_buffer_pool *buffer_pool);
 
 AWS_EXTERN_C_END
 

include/aws/s3/private/s3_meta_request_impl.h

@@ -170,6 +170,7 @@ struct aws_s3_meta_request {
     /* Customer specified callbacks. */
     aws_s3_meta_request_headers_callback_fn *headers_callback;
     aws_s3_meta_request_receive_body_callback_fn *body_callback;
+    aws_s3_meta_request_receive_body_callback_ex_fn *body_callback_ex;
     aws_s3_meta_request_finish_fn *finish_callback;
     aws_s3_meta_request_shutdown_fn *shutdown_callback;
     aws_s3_meta_request_progress_fn *progress_callback;
@@ -226,8 +227,11 @@ struct aws_s3_meta_request {
         /* To track aws_s3_requests with cancellable HTTP streams */
         struct aws_linked_list cancellable_http_streams_list;
 
+        /* To track aws_future_s3_buffers that might need to be cleaned up on cancel */
+        struct aws_linked_list pending_buffer_futures;
+
         /* Data for async-writes. */
-        struct {
+        struct aws_s3_async_write_data {
             /* Whether a part request can be sent (we have 1 part's worth of data, or EOF) */
             bool ready_to_send;
 
@@ -237,7 +241,8 @@ struct aws_s3_meta_request {
             /* Holds buffered data we can't immediately send.
              * The length will always be less than part-size */
             struct aws_byte_buf buffered_data;
-            struct aws_s3_buffer_pool_ticket *buffered_data_ticket;
+            struct aws_s3_buffer_ticket *buffered_data_ticket;
+            struct aws_future_s3_buffer_ticket *buffered_ticket_future;
 
             /* Waker callback.
              * Stored if a poll_write() call returns result.is_pending
@@ -401,6 +406,9 @@ bool aws_s3_meta_request_are_events_out_for_delivery_synced(struct aws_s3_meta_r
 /* Cancel the requests with cancellable HTTP stream for the meta request */
 void aws_s3_meta_request_cancel_cancellable_requests_synced(struct aws_s3_meta_request *meta_request, int error_code);
 
+/* Cancel the pending buffer futures for the meta request */
+void aws_s3_meta_request_cancel_pending_buffer_futures_synced(struct aws_s3_meta_request *meta_request, int error_code);
+
 /* Asynchronously read from the meta request's input stream. Should always be done outside of any mutex,
  * as reading from the stream could cause user code to call back into aws-c-s3.
  * This will fill the buffer to capacity, unless end of stream is reached.

include/aws/s3/private/s3_request.h

@@ -12,7 +12,6 @@
 #include <aws/common/thread.h>
 #include <aws/s3/s3.h>
 
-#include <aws/s3/private/s3_buffer_pool.h>
 #include <aws/s3/private/s3_checksums.h>
 
 struct aws_http_message;
@@ -21,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,
 };
 
 /**
@@ -116,10 +114,16 @@ struct aws_s3_request {
     /* Linked list node used for tracking the request is active from HTTP level. */
     struct aws_linked_list_node cancellable_http_streams_list_node;
 
+    /* Linked list node used for tracking buffer acquire futures. */
+    struct aws_linked_list_node pending_buffer_future_list_node;
+
     /* The meta request lock must be held to access the data */
     struct {
         /* The underlying http stream, only valid when the request is active from HTTP level */
         struct aws_http_stream *cancellable_http_stream;
+
+        /* Buffer future. */
+        struct aws_future_s3_buffer_ticket *buffer_future;
     } synced_data;
 
     /* TODO Ref count on the request is no longer needed--only one part of code should ever be holding onto a request,
@@ -135,7 +139,7 @@ struct aws_s3_request {
      * retried.*/
     struct aws_byte_buf request_body;
 
-    struct aws_s3_buffer_pool_ticket *ticket;
+    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. */
@@ -225,11 +229,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;

include/aws/s3/s3.h

@@ -49,6 +49,7 @@ enum aws_s3_errors {
     AWS_ERROR_S3_RECV_FILE_ALREADY_EXISTS,
     AWS_ERROR_S3_RECV_FILE_NOT_FOUND,
     AWS_ERROR_S3_REQUEST_TIMEOUT,
+    AWS_ERROR_S3_BUFFER_ALLOCATION_FAILED,
 
     AWS_ERROR_S3_END_RANGE = AWS_ERROR_ENUM_END_RANGE(AWS_C_S3_PACKAGE_ID)
 };

include/aws/s3/s3_buffer_pool.h

@@ -0,0 +1,146 @@
+#ifndef AWS_S3_BUFFER_POOL_H
+#define AWS_S3_BUFFER_POOL_H
+
+#include <aws/common/ref_count.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.
+ */
+
+/**
+ * Generic memory pool interface.
+ * Allows consumers of aws-c-s3 to override how buffer allocation for part buffers is done.
+ * Refer to docs/memory_aware_request_execution.md for details on how default implementation works.
+ * WARNING: this is currently experimental feature and does not provide API stability guarantees so should be used with
+ * caution. At highlevel the flow is as follows:
+ * - crt scheduler queues up requests to be prepared
+ * - requests being prepared will try to reserve mem (i.e. obtain a ticket) and wait until they get it before proceeding
+ * - once mem is reserved requests will proceed with the pipeline
+ * - request will acquire buffer from the ticket when needed
+ * - ticket is released when request is done with the buffer
+ * Note: in some cases pipeline can stall if new buffer cannot be allocated (ex. async writes flow).
+ * In this case reserve request will indicate that not granting the ticket can block and buffer pool should try to
+ * allocate ticket right away (or wait and call waker when mem is allocated for the case of async writes).
+ */
+
+AWS_PUSH_SANE_WARNING_LEVEL
+AWS_EXTERN_C_BEGIN
+struct aws_s3_buffer_ticket;
+
+/**
+ * aws_future<aws_s3_buffer_ticket*>
+ * Buffer ticket future used for reservations.
+ */
+AWS_FUTURE_T_POINTER_WITH_RELEASE_DECLARATION(aws_future_s3_buffer_ticket, struct aws_s3_buffer_ticket, AWS_S3_API)
+
+/**
+ * Meta information about ticket reservation request.
+ */
+struct aws_s3_buffer_pool_reserve_meta {
+    /* client reserving the ticket. accounts for buffer pool being shared between clients. */
+    struct aws_s3_client *client;
+
+    /* meta request ticket is being reserved for. */
+    struct aws_s3_meta_request *meta_request;
+
+    /* size of the buffer to reserve. */
+    size_t size;
+
+    /* whether not granting reservation can result in request pipeline being blocked.
+     * Note: blocking is currently a terminal condition and that cannot be recovered from,
+     * i.e. meta request will be stuck and not make any process.
+     * As such buffer pool should either grant or error out reservation in sync.
+     * This scenario currently only occurs in the async_write flows. */
+    bool can_block;
+};
+
+struct aws_s3_buffer_ticket;
+
+struct aws_s3_buffer_ticket_vtable {
+    /**
+     * Get buffer associated with the ticket.
+     * Note: can be called multiple times and the same buffer should be returned. In some cases ticket might not be
+     * claimed at all.
+     */
+    struct aws_byte_buf (*claim)(struct aws_s3_buffer_ticket *ticket);
+
+    /* Implement below for custom ref count behavior. Alternatively set those to null and init the ref count. */
+    struct aws_s3_buffer_ticket *(*acquire)(struct aws_s3_buffer_ticket *ticket);
+    struct aws_s3_buffer_ticket *(*release)(struct aws_s3_buffer_ticket *ticket);
+};
+
+/**
+ * Polymorphic ticket.
+ */
+struct aws_s3_buffer_ticket {
+    struct aws_s3_buffer_ticket_vtable *vtable;
+    struct aws_ref_count ref_count;
+    void *impl;
+};
+
+AWS_S3_API struct aws_byte_buf aws_s3_buffer_ticket_claim(struct aws_s3_buffer_ticket *ticket);
+
+AWS_S3_API struct aws_s3_buffer_ticket *aws_s3_buffer_ticket_acquire(struct aws_s3_buffer_ticket *ticket);
+AWS_S3_API struct aws_s3_buffer_ticket *aws_s3_buffer_ticket_release(struct aws_s3_buffer_ticket *ticket);
+
+struct aws_s3_buffer_pool;
+
+struct aws_s3_buffer_pool_vtable {
+    /* Reserve a ticket. Returns a future that is granted whenever reservation can be made. */
+    struct aws_future_s3_buffer_ticket *(
+        *reserve)(struct aws_s3_buffer_pool *pool, struct aws_s3_buffer_pool_reserve_meta meta);
+
+    /**
+     * Trim the pool. This is mostly a suggestion, which pool can decide to ignore. Triggered by CRT when
+     * client has been idle for some time.
+     **/
+    void (*trim)(struct aws_s3_buffer_pool *pool);
+
+    /* Implement below for custom ref count behavior. Alternatively set those to null and init the ref count. */
+    struct aws_s3_buffer_pool *(*acquire)(struct aws_s3_buffer_pool *pool);
+    struct aws_s3_buffer_pool *(*release)(struct aws_s3_buffer_pool *pool);
+};
+
+/**
+ * Polymorphic buffer pool.
+ */
+struct aws_s3_buffer_pool {
+    struct aws_s3_buffer_pool_vtable *vtable;
+    struct aws_ref_count ref_count;
+    void *impl;
+};
+
+AWS_S3_API struct aws_future_s3_buffer_ticket *aws_s3_buffer_pool_reserve(
+    struct aws_s3_buffer_pool *buffer_pool,
+    struct aws_s3_buffer_pool_reserve_meta meta);
+AWS_S3_API void aws_s3_buffer_pool_trim(struct aws_s3_buffer_pool *buffer_pool);
+
+AWS_S3_API struct aws_s3_buffer_pool *aws_s3_buffer_pool_acquire(struct aws_s3_buffer_pool *buffer_pool);
+AWS_S3_API struct aws_s3_buffer_pool *aws_s3_buffer_pool_release(struct aws_s3_buffer_pool *buffer_pool);
+
+/**
+ * Buffer pool configuration options.
+ */
+struct aws_s3_buffer_pool_config {
+    struct aws_s3_client *client; /* Client creating the pool. */
+    size_t part_size;             /* Default part size of the client. */
+    size_t max_part_size;         /* Max part size configured on the client. */
+    size_t memory_limit;          /* Memory limit set on the client. */
+};
+
+/**
+ * Factory to construct the pool for the given config. Passes along buffer related info configured on the client, which
+ * factory may ignore when considering how to construct pool.
+ * This implementation should fail if pool cannot be constructed for some reason (ex. if config params cannot be met),
+ * by logging failure reason, returning null and raising aws_error.
+ */
+typedef struct aws_s3_buffer_pool *(aws_s3_buffer_pool_factory_fn)(struct aws_allocator *allocator,
+                                                                   struct aws_s3_buffer_pool_config config);
+
+AWS_EXTERN_C_END
+AWS_POP_SANE_WARNING_LEVEL
+
+#endif /* AWS_S3_BUFFER_POOL_H */