diff --git a/include/aws/s3/private/s3_request.h b/include/aws/s3/private/s3_request.h index 33776b001..393ecc2b5 100644 --- a/include/aws/s3/private/s3_request.h +++ b/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 diff --git a/source/s3_auto_ranged_get.c b/source/s3_auto_ranged_get.c index 5a3249531..7b4f3178f 100644 --- a/source/s3_auto_ranged_get.c +++ b/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, diff --git a/source/s3_auto_ranged_put.c b/source/s3_auto_ranged_put.c index 0c11315fe..be5d7a352 100644 --- a/source/s3_auto_ranged_put.c +++ b/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, diff --git a/source/s3_client.c b/source/s3_client.c index be97b9b88..16cc09192 100644 --- a/source/s3_client.c +++ b/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 = diff --git a/source/s3_meta_request.c b/source/s3_meta_request.c index 50d856661..f1a86a39b 100644 --- a/source/s3_meta_request.c +++ b/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; diff --git a/source/s3_request.c b/source/s3_request.c index fa7cd83ba..179a6b56b 100644 --- a/source/s3_request.c +++ b/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);