Dynamic default part size (#575)

Co-authored-by: Dmitriy Musatkin <63878209+DmitriyMusatkin@users.noreply.github.com>

Author: TingDaoK

.github/workflows/ci.yml

@@ -223,6 +223,8 @@ jobs:
       uses: actions/checkout@v4
     - name: Build ${{ env.PACKAGE_NAME }} + consumers
       run: |
+        python3 -m venv .venv
+        source .venv/bin/activate
         python3 -c "from urllib.request import urlretrieve; urlretrieve('${{ env.BUILDER_HOST }}/${{ env.BUILDER_SOURCE }}/${{ env.BUILDER_VERSION }}/builder.pyz?run=${{ env.RUN }}', 'builder')"
         chmod a+x builder
         ./builder build -p ${{ env.PACKAGE_NAME }} --cmake-extra=-DASSERT_LOCK_HELD=ON

CMakeLists.txt

@@ -108,6 +108,7 @@ install(FILES "${CMAKE_CURRENT_BINARY_DIR}/${PROJECT_NAME}-config.cmake"
 
 include(CTest)
 if (BUILD_TESTING)
+    add_definitions(-DAWS_C_S3_ENABLE_TEST_STUBS)
     add_subdirectory(tests)
     if (NOT BYO_CRYPTO AND NOT CMAKE_CROSSCOMPILING)
         add_subdirectory(samples)

include/aws/s3/private/s3_auto_ranged_get.h

@@ -21,6 +21,18 @@ struct aws_s3_auto_ranged_get {
 
     struct aws_string *etag;
 
+    /* Estimated object stored part size based on ETag analysis */
+    uint64_t estimated_object_stored_part_size;
+    /* Number of parts stored in S3. We derive this from ETag, if ETag is not formatted as expected, this will be
+     * default to 1.
+     * Note: For S3Express Append, the object will be treated as a single part, even though, it can be multiple parts
+     * stored in S3.
+     */
+    uint64_t num_stored_parts;
+    /* Part size was set or not from user for this meta request. */
+    bool part_size_set;
+    bool force_dynamic_part_size;
+
     bool initial_message_has_start_range;
     bool initial_message_has_end_range;
     uint64_t initial_range_start;
@@ -74,6 +86,7 @@ AWS_S3_API struct aws_s3_meta_request *aws_s3_meta_request_auto_ranged_get_new(
     struct aws_allocator *allocator,
     struct aws_s3_client *client,
     size_t part_size,
+    bool part_size_set,
     const struct aws_s3_meta_request_options *options);
 
 AWS_EXTERN_C_END

include/aws/s3/private/s3_client_impl.h

@@ -178,7 +178,10 @@ struct aws_s3_client_vtable {
         struct aws_http_connection *client_connection,
         const struct aws_http_make_request_options *options);
 
-    void (*after_prepare_upload_part_finish)(struct aws_s3_request *request, struct aws_http_message *message);
+#ifdef AWS_C_S3_ENABLE_TEST_STUBS
+    /********************* TEST ONLY STUB **************************/
+    void (*after_prepare_upload_part_finish_stub)(struct aws_s3_request *request, struct aws_http_message *message);
+#endif
 };
 
 struct aws_s3_upload_part_timeout_stats {
@@ -234,10 +237,17 @@ struct aws_s3_client {
      * to meta requests for use. */
     const size_t part_size;
 
+    bool part_size_set;
+
     /* Size of parts for files when doing gets or puts.  This exists on the client as configurable option that is passed
      * to meta requests for use. */
     const uint64_t max_part_size;
 
+    /* Calculated optimal range size for GET operations based on client configuration (memory limits, throughput
+     * targets). This is used when part_size is not explicitly configured, replacing the default with reasonable
+     * calculation. Value is calculated during client initialization and remains constant for the client's lifetime. */
+    const uint64_t optimal_range_size;
+
     /* File I/O options. */
     bool fio_options_set;
     struct aws_s3_file_io_options fio_opts;

include/aws/s3/private/s3_default_buffer_pool.h

@@ -6,6 +6,8 @@
  * SPDX-License-Identifier: Apache-2.0.
  */
 
+#include <aws/common/hash_table.h>
+#include <aws/common/mutex.h>
 #include <aws/s3/s3.h>
 #include <aws/s3/s3_buffer_pool.h>
 
@@ -59,11 +61,79 @@ struct aws_s3_default_buffer_pool_usage_stats {
     /* Secondary memory reserved, but not yet used. Accurate, maps directly to base allocator. */
     size_t secondary_reserved;
 
+    /* Overall memory allocated for special-sized blocks. */
+    size_t special_blocks_allocated;
+    /* Number of special block sizes created. */
+    size_t special_blocks_num;
+    /* Memory reserved in special-sized blocks. */
+    size_t special_blocks_reserved;
+    /* Memory used in special-sized blocks. */
+    size_t special_blocks_used;
+
     /* Bytes used in "forced" buffers (created even if they exceed memory limits).
      * This is always <= primary_used + secondary_used */
     size_t forced_used;
 };
 
+/* Structure to track special-sized blocks */
+struct s3_special_block_list {
+    struct aws_allocator *allocator;
+    size_t buffer_size;           /* Size of buffers in this list */
+    struct aws_array_list blocks; /* Array of uint8_t* pointers to allocated blocks */
+};
+
+struct aws_s3_default_buffer_pool {
+    struct aws_allocator *base_allocator;
+    struct aws_mutex mutex;
+
+    size_t block_size;
+    size_t chunk_size;
+    /* size at which allocations should go to secondary */
+    size_t primary_size_cutoff;
+
+    /* NOTE: See aws_s3_buffer_pool_usage_stats for descriptions of most fields */
+
+    size_t mem_limit;
+
+    size_t primary_allocated;
+    size_t primary_reserved;
+    size_t primary_used;
+
+    size_t special_blocks_allocated;
+    size_t special_blocks_reserved;
+    size_t special_blocks_used;
+
+    size_t secondary_reserved;
+    size_t secondary_used;
+
+    size_t forced_used;
+
+    struct aws_array_list blocks;
+
+    struct aws_linked_list pending_reserves;
+
+    /* Special-sized blocks: hash table mapping size -> struct s3_special_block_list * */
+    /* TODO: let's discuss about the special list lifetime. Should we just keep it with the memory pool? Concern is that
+     * the pool will live with the client, and may result in all sorts of special lists to be around. */
+    struct aws_hash_table special_blocks;
+
+    /* TEST ONLY: to force the special blocks alive during trim. */
+    bool force_keeping_special_blocks;
+};
+
+struct s3_pending_reserve {
+    struct aws_linked_list_node node;
+    struct aws_future_s3_buffer_ticket *ticket_future;
+    struct aws_s3_default_buffer_ticket *ticket;
+    struct aws_s3_buffer_pool_reserve_meta meta;
+};
+
+struct s3_buffer_pool_block {
+    size_t block_size;
+    uint8_t *block_ptr;
+    uint16_t alloc_bit_mask;
+};
+
 /*
  * Create new buffer pool.
  * chunk_size - specifies the size of memory that will most commonly be acquired

include/aws/s3/private/s3_meta_request_impl.h

@@ -122,6 +122,12 @@ struct aws_s3_meta_request_vtable {
 
     /* Pause the given request */
     int (*pause)(struct aws_s3_meta_request *meta_request, struct aws_s3_meta_request_resume_token **resume_token);
+
+#ifdef AWS_C_S3_ENABLE_TEST_STUBS
+    /********************* TEST ONLY STUB **************************/
+    /* A stub to the update implementation from meta request with the lock held. Only for tests. */
+    bool (*synced_update_stub)(struct aws_s3_meta_request *meta_request);
+#endif
 };
 
 /**
@@ -151,6 +157,8 @@ struct aws_s3_meta_request {
 
     /* Part size to use for uploads and downloads.  Passed down by the creating client. */
     const size_t part_size;
+    /* Hard limit on max connections set through the meta request option. */
+    const uint32_t max_active_connections_override;
 
     struct aws_cached_signing_config_aws *cached_signing_config;
 
@@ -160,6 +168,9 @@ struct aws_s3_meta_request {
 
     struct aws_s3_endpoint *endpoint;
 
+    /* Number of requests being sent/received over network for the meta request. */
+    struct aws_atomic_var num_requests_network;
+
     /* Event loop to schedule IO work related on, ie, reading from streams, streaming parts back to the caller, etc...
      * After the meta request is finished, this will be reset along with the client reference.*/
     struct aws_event_loop *io_event_loop;
@@ -179,6 +190,10 @@ struct aws_s3_meta_request {
 
     enum aws_s3_meta_request_type type;
     struct aws_string *s3express_session_host;
+    /* Is the meta request made to s3express bucket or not. */
+    bool is_express;
+    /* If the buffer pool optimized for the specific size or not. */
+    bool buffer_pool_optimized;
 
     struct {
         struct aws_mutex lock;
@@ -263,6 +278,9 @@ struct aws_s3_meta_request {
         /* True if this meta request is currently in the client's list. */
         bool scheduled;
 
+        /* Track the number of requests being prepared for this meta request. */
+        size_t num_request_being_prepared;
+
     } client_process_work_threaded_data;
 
     /* Anything in this structure should only ever be accessed by the meta-request from its io_event_loop thread. */
@@ -271,6 +289,9 @@ struct aws_s3_meta_request {
          * This is an optimization, we could have just copied the array when the task runs,
          * but swapping two array-lists back and forth avoids an allocation. */
         struct aws_array_list event_delivery_array;
+
+        /* The range start for the next response body delivery */
+        uint64_t next_deliver_range_start;
     } io_threaded_data;
 
     const bool should_compute_content_md5;
@@ -407,9 +428,11 @@ void aws_s3_meta_request_add_event_for_delivery_synced(
 bool aws_s3_meta_request_are_events_out_for_delivery_synced(struct aws_s3_meta_request *meta_request);
 
 /* Cancel the requests with cancellable HTTP stream for the meta request */
+AWS_S3_API
 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 */
+AWS_S3_API
 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,

include/aws/s3/private/s3_request.h

@@ -168,8 +168,20 @@ struct aws_s3_request_metrics {
         int error_code;
         /* Retry attempt. */
         uint32_t retry_attempt;
+        /* Is the memory for the request allocated from the buffer pool or not. */
+        bool memory_allocated_from_pool;
     } crt_info_metrics;
 
+    /* TODO: align the part info metrics with the others, eg: `aws_s3_mpu_part_info`. */
+    struct {
+        /* Beginning range of this part. */
+        uint64_t part_range_start;
+        /* Last byte of this part. */
+        uint64_t part_range_end;
+        /* Part number that this request refers to. */
+        uint32_t part_number;
+    } part_info_metrics;
+
     struct aws_ref_count ref_count;
 };
 
@@ -229,11 +241,9 @@ 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. */
     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. */
     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.)

include/aws/s3/private/s3_util.h

@@ -160,6 +160,23 @@ extern const double g_default_throughput_target_gbps;
 
 AWS_S3_API
 extern const uint64_t g_streaming_object_size_threshold;
+
+AWS_S3_API
+extern const uint64_t g_default_part_size_fallback;
+
+AWS_S3_API
+extern const uint64_t g_default_max_part_size;
+
+AWS_S3_API
+extern const uint64_t g_s3_optimal_range_size_alignment;
+
+AWS_S3_API
+extern const uint32_t g_s3express_connection_limitation;
+AWS_S3_API
+extern const uint64_t g_s3express_connection_limitation_part_size_threshold;
+AWS_S3_API
+extern const uint64_t g_s3express_connection_limitation_object_size_threshold;
+
 /**
  * Returns AWS_S3_REQUEST_TYPE_UNKNOWN if name doesn't map to an enum value.
  */
@@ -240,12 +257,20 @@ void aws_s3_add_user_agent_header(struct aws_allocator *allocator, struct aws_ht
  * object-size. All output arguments are optional.*/
 AWS_S3_API
 int aws_s3_parse_content_range_response_header(
-    struct aws_allocator *allocator,
     struct aws_http_headers *response_headers,
     uint64_t *out_range_start,
     uint64_t *out_range_end,
     uint64_t *out_object_size);
 
+/* Given a Content-Range header value as a byte cursor, parses the range-start, range-end and
+ * object-size. All output arguments are optional. */
+AWS_S3_API
+int aws_s3_parse_content_range_cursor(
+    struct aws_byte_cursor content_range_cursor,
+    uint64_t *out_range_start,
+    uint64_t *out_range_end,
+    uint64_t *out_object_size);
+
 /* Given response headers, parses the content-length from a content-length response header.*/
 AWS_S3_API
 int aws_s3_parse_content_length_response_header(
@@ -318,6 +343,53 @@ int aws_s3_check_headers_for_checksum(
     struct aws_byte_buf *out_checksum_buffer,
     bool meta_request_level);
 
+/**
+ * Calculate client-level optimal range size based on memory and connection constraints.
+ * This function is called during client initialization to determine the base range size
+ * using the formula: MemoryLimit / concurrency / divisor.
+ * The result is rounded up to ensure proper alignment and applies minimum size constraints.
+ *
+ * @param memory_limit_in_bytes Total memory limit available for buffering
+ * @param max_connections Maximum number of concurrent connections
+ * @param out_client_optimal_range_size Output parameter for calculated client-level optimal range size
+ * @return AWS_OP_SUCCESS on success, AWS_OP_ERR on failure (caller should fall back to default)
+ */
+AWS_S3_API
+int aws_s3_calculate_client_optimal_range_size(
+    uint64_t memory_limit_in_bytes,
+    uint32_t max_connections,
+    uint64_t *out_client_optimal_range_size);
+
+/**
+ * Calculate request-level optimal range size by considering object-specific information.
+ * This function is called per request to adjust the client-level range size based on
+ * estimated object stored part size using: min(client_optimal_range_size, estimated_object_stored_part_size).
+ *
+ * @param client_optimal_range_size The client-level optimal range size from initialization
+ * @param estimated_object_stored_part_size Estimated size of object stored parts in S3
+ * @param is_express If the request is a s3express request or not.
+ * @param out_request_optimal_range_size Output parameter for calculated request-level optimal range size
+ * @return AWS_OP_SUCCESS on success, AWS_OP_ERR on failure (caller should fall back to client size)
+ */
+AWS_S3_API
+int aws_s3_calculate_request_optimal_range_size(
+    uint64_t client_optimal_range_size,
+    uint64_t estimated_object_stored_part_size,
+    bool is_express,
+    uint64_t *out_request_optimal_range_size);
+
+/**
+ * Extract the number of parts from an S3 ETag header value.
+ * S3 multipart upload ETags have the format "<hash>-<number_of_parts>".
+ * Single-part uploads have ETags without dashes.
+ *
+ * @param etag_header_value The ETag header value (may include quotes)
+ * @param out_num_parts Output parameter for the number of parts (1 for single-part uploads)
+ * @return AWS_OP_SUCCESS on success, AWS_OP_ERR on failure (invalid ETag format)
+ */
+AWS_S3_API
+int aws_s3_extract_parts_from_etag(struct aws_byte_cursor etag_header_value, uint32_t *out_num_parts);
+
 AWS_EXTERN_C_END
 
 #endif /* AWS_S3_UTIL_H */