@@ -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 */
9696static 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+
601617struct 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
0 commit comments