Skip to content

Commit 3df068d

Browse files
extra tests
1 parent 439e3d5 commit 3df068d

3 files changed

Lines changed: 55 additions & 20 deletions

File tree

source/s3_auto_ranged_get.c

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -301,8 +301,7 @@ static bool s_s3_auto_ranged_get_update(
301301
AWS_LOGF_INFO(
302302
AWS_LS_S3_META_REQUEST,
303303
"id=%p: Doing a ranged get to discover the size of the object and get the first part",
304-
(void *)meta_request
305-
);
304+
(void *)meta_request);
306305

307306
request = aws_s3_request_new(
308307
meta_request,

source/s3_default_buffer_pool.c

Lines changed: 45 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ struct aws_s3_default_buffer_ticket {
5757
size_t chunks_used;
5858
bool forced;
5959
struct aws_s3_buffer_pool *pool;
60-
bool is_special_block; /* True if this ticket is from a special-sized block */
60+
bool is_special_block; /* True if this ticket is from a special-sized block */
6161
enum aws_s3_default_buffer_pool_reserved_from reserved_from; /* which area ticket was reserved from */
6262
};
6363

@@ -91,7 +91,7 @@ static const size_t s_max_impact_of_forced_buffers_on_memory_limit_as_percentage
9191

9292
/*
9393
* primary min cut off. below this number buffers are not allocated from primary and instead use
94-
* secondary.
94+
* secondary.
9595
*/
9696
static const size_t s_primary_min_cutoff = KB_TO_BYTES(128);
9797

@@ -246,10 +246,10 @@ struct aws_s3_buffer_pool *aws_s3_default_buffer_pool_new(
246246
}
247247

248248
/*
249-
* This is fairly arbitrary. System allocator tends to be a bit better on tiny allocations.
250-
* And we want to avoid wasting mem in that case.
251-
* Note this is not the answer for wasting mem, but it does improve some corner cases.
252-
*/
249+
* This is fairly arbitrary. System allocator tends to be a bit better on tiny allocations.
250+
* And we want to avoid wasting mem in that case.
251+
* Note this is not the answer for wasting mem, but it does improve some corner cases.
252+
*/
253253
size_t min_primary_size = 0;
254254
if (4 * chunk_size >= s_primary_min_cutoff) {
255255
min_primary_size = s_primary_min_cutoff;
@@ -598,6 +598,22 @@ static bool s_should_trim_for_reserve_synced(
598598
return true;
599599
}
600600

601+
static size_t s_align_up_to_block_size(size_t size, size_t block_size, size_t *out) {
602+
AWS_FATAL_ASSERT(block_size > 0);
603+
size_t remainder = size % block_size;
604+
if (remainder == 0)
605+
return size;
606+
607+
size_t sub = 0;
608+
if (!aws_sub_size_checked(block_size, remainder, &sub)) {
609+
if (!aws_add_size_checked(size, sub, out)) {
610+
return AWS_OP_SUCCESS;
611+
}
612+
}
613+
614+
return AWS_OP_ERR;
615+
}
616+
601617
struct aws_s3_default_buffer_ticket *s_try_reserve_synced(
602618
struct aws_s3_buffer_pool *buffer_pool_wrapper,
603619
struct aws_s3_buffer_pool_reserve_meta meta) {
@@ -641,13 +657,30 @@ struct aws_s3_default_buffer_ticket *s_try_reserve_synced(
641657
ticket->is_special_block = true;
642658
ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_SPECIAL;
643659
buffer_pool->special_blocks_reserved += meta.size;
644-
} else if (meta.size <= buffer_pool->primary_size_cutoff) { /* TODO: this needs to be smarter. what if block size more than mem_lim*/
660+
} else if (meta.size <= buffer_pool->primary_size_cutoff) {
661+
/* This needs to be smarter. Currently if primary req size is below limit, it will allocate full block,
662+
which can be above limit. */
645663
if (meta.size <= buffer_pool->primary_size_min_cutoff) {
646664
ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY;
647665
buffer_pool->secondary_reserved += meta.size;
648666
} else {
649-
ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_PRIMARY;
650-
buffer_pool->primary_reserved += meta.size;
667+
size_t aligned = 0;
668+
if (s_align_up_to_block_size(
669+
buffer_pool->primary_reserved + meta.size, buffer_pool->block_size, &aligned)) {
670+
/* aligning failed for some reason, just revert to secondary */
671+
ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY;
672+
buffer_pool->secondary_reserved += meta.size;
673+
} else {
674+
if ((overall_taken + aligned) > buffer_pool->mem_limit) {
675+
/* allocating this from primary would result in exceeding mem limit if new block is created,
676+
so allocate from secondary instead. */
677+
ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY;
678+
buffer_pool->secondary_reserved += meta.size;
679+
} else {
680+
ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_PRIMARY;
681+
buffer_pool->primary_reserved += meta.size;
682+
}
683+
}
651684
}
652685
} else {
653686
ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY;
@@ -785,7 +818,7 @@ static struct aws_byte_buf s_acquire_buffer_synced(
785818
ticket->is_special_block);
786819

787820
/* Check if this is a special-sized allocation */
788-
if (ticket->is_special_block) {
821+
if (ticket->reserved_from == AWS_S3_BUFFER_POOL_RESERVED_FROM_SPECIAL) {
789822
struct aws_hash_element *elem = NULL;
790823
aws_hash_table_find(&buffer_pool->special_blocks, (void *)ticket->size, &elem);
791824
AWS_FATAL_ASSERT(elem != NULL);
@@ -826,9 +859,10 @@ static struct aws_byte_buf s_acquire_buffer_synced(
826859
}
827860
buffer_pool->special_blocks_used += ticket->size;
828861
buffer_pool->special_blocks_reserved -= ticket->size;
829-
} else if (ticket->size <= buffer_pool->primary_size_cutoff) {
862+
} else if (ticket->reserved_from == AWS_S3_BUFFER_POOL_RESERVED_FROM_PRIMARY) {
830863
ticket->ptr = s_primary_acquire_synced(buffer_pool, ticket);
831864
} else {
865+
AWS_ASSERT(ticket->reserved_from == AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY)
832866
ticket->ptr = aws_mem_acquire(buffer_pool->base_allocator, ticket->size);
833867
buffer_pool->secondary_used += ticket->size;
834868

tests/s3_default_buffer_pool_tests.c

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,7 @@ static int s_test_s3_buffer_pool_limits(struct aws_allocator *allocator, void *c
139139

140140
struct aws_future_s3_buffer_ticket *future2 = aws_s3_default_buffer_pool_reserve(
141141
buffer_pool, (struct aws_s3_buffer_pool_reserve_meta){.size = MB_TO_BYTES(32)});
142+
142143
ASSERT_NOT_NULL(future2);
143144
ASSERT_TRUE(aws_future_s3_buffer_ticket_is_done(future2));
144145
ASSERT_INT_EQUALS(aws_future_s3_buffer_ticket_get_error(future2), AWS_OP_SUCCESS);
@@ -147,6 +148,9 @@ static int s_test_s3_buffer_pool_limits(struct aws_allocator *allocator, void *c
147148
struct aws_byte_buf buf2 = aws_s3_buffer_ticket_claim(ticket2);
148149
ASSERT_NOT_NULL(buf2.buffer);
149150

151+
struct aws_s3_default_buffer_pool_usage_stats stats = aws_s3_default_buffer_pool_get_usage(buffer_pool);
152+
ASSERT_UINT_EQUALS(0, stats.primary_allocated);
153+
150154
for (size_t i = 0; i < 6; ++i) {
151155
aws_s3_buffer_ticket_release(tickets[i]);
152156
aws_future_s3_buffer_ticket_release(ticket_futures[i]);
@@ -636,12 +640,11 @@ AWS_TEST_CASE(
636640
test_s3_buffer_pool_forced_buffer_wont_stop_reservations,
637641
s_test_s3_buffer_pool_forced_buffer_wont_stop_reservations)
638642

639-
640643
/* Test small requests reserve expected amount of mem without rounding up to part size. */
641644
static int s_test_s3_buffer_pool_reserve_tiny_chunks(struct aws_allocator *allocator, void *ctx) {
642645
(void)ctx;
643646
const size_t chunk_size = MB_TO_BYTES(8);
644-
const size_t small_size = MB_TO_BYTES(1);
647+
const size_t small_size = KB_TO_BYTES(1);
645648
const size_t mem_limit = GB_TO_BYTES(1);
646649
struct aws_s3_buffer_pool *buffer_pool = aws_s3_default_buffer_pool_new(
647650
allocator, (struct aws_s3_buffer_pool_config){.part_size = chunk_size, .memory_limit = mem_limit});
@@ -651,22 +654,21 @@ static int s_test_s3_buffer_pool_reserve_tiny_chunks(struct aws_allocator *alloc
651654
aws_s3_default_buffer_pool_reserve(buffer_pool, (struct aws_s3_buffer_pool_reserve_meta){.size = small_size});
652655
ASSERT_TRUE(aws_future_s3_buffer_ticket_is_done(future));
653656
ASSERT_INT_EQUALS(aws_future_s3_buffer_ticket_get_error(future), AWS_OP_SUCCESS);
654-
normal_ticket = aws_future_s3_buffer_ticket_get_result_by_move(future);
657+
ticket = aws_future_s3_buffer_ticket_get_result_by_move(future);
655658
struct aws_byte_buf buf = aws_s3_buffer_ticket_claim(ticket);
656-
ASSERT_UINT_EQUALS(small, buf.capacity);
659+
ASSERT_UINT_EQUALS(small_size, buf.capacity);
657660

658661
struct aws_s3_default_buffer_pool_usage_stats stats = aws_s3_default_buffer_pool_get_usage(buffer_pool);
659662
ASSERT_INT_EQUALS(small_size, stats.secondary_used);
660663

661664
aws_s3_buffer_ticket_release(ticket);
662665
aws_future_s3_buffer_ticket_release(future);
663666

667+
stats = aws_s3_default_buffer_pool_get_usage(buffer_pool);
664668
ASSERT_INT_EQUALS(0, stats.secondary_used);
665669

666670
/* Cleanup */
667671
aws_s3_default_buffer_pool_destroy(buffer_pool);
668672
return 0;
669673
}
670-
AWS_TEST_CASE(
671-
test_s3_buffer_pool_reserve_tiny_chunks,
672-
s_test_s3_buffer_pool_reserve_tiny_chunks)
674+
AWS_TEST_CASE(test_s3_buffer_pool_reserve_tiny_chunks, s_test_s3_buffer_pool_reserve_tiny_chunks)

0 commit comments

Comments
 (0)