update logic

Author: DmitriyMusatkin

include/aws/s3/private/s3_default_buffer_pool.h

@@ -91,6 +91,9 @@ struct aws_s3_default_buffer_pool {
     /* size at which allocations should go to secondary */
     size_t primary_size_cutoff;
 
+    /* size below which allocations should go to secondary */
+    size_t primary_size_min_cutoff;
+
     /* NOTE: See aws_s3_buffer_pool_usage_stats for descriptions of most fields */
 
     size_t mem_limit;

source/s3_default_buffer_pool.c

@@ -45,13 +45,20 @@
 
 struct aws_s3_default_buffer_pool;
 
+enum aws_s3_default_buffer_pool_reserved_from {
+    AWS_S3_BUFFER_POOL_RESERVED_FROM_PRIMARY,
+    AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY,
+    AWS_S3_BUFFER_POOL_RESERVED_FROM_SPECIAL,
+};
+
 struct aws_s3_default_buffer_ticket {
     size_t size;
     uint8_t *ptr;
     size_t chunks_used;
     bool forced;
     struct aws_s3_buffer_pool *pool;
     bool is_special_block; /* True if this ticket is from a special-sized block */
+    enum aws_s3_default_buffer_pool_reserved_from reserved_from /* which area ticket was reserved from */
 };
 
 /* Default size for blocks array. Note: this is just for meta info, blocks
@@ -82,6 +89,12 @@ static const size_t s_max_chunk_size_for_buffer_reuse = MB_TO_BYTES(64);
  * we still consider 1GiB available for normal buffer usage. */
 static const size_t s_max_impact_of_forced_buffers_on_memory_limit_as_percentage = 80;
 
+/*
+ * primary min cut off. below this number buffers are not allocated from primary and instead use
+ * secondary. 
+ */
+static const size_t s_primary_min_cutoff = KB_TO_BYTES(128);
+
 /*
  * Sets n bits at position starting with LSB.
  * Note: n must be at most 8, but in practice will always be at most 4.
@@ -232,6 +245,18 @@ struct aws_s3_buffer_pool *aws_s3_default_buffer_pool_new(
         chunk_size = 0;
     }
 
+    /*
+    * This is fairly arbitrary. System allocator tends to be a bit better on tiny allocations.
+    * And we want to avoid wasting mem in that case.
+    * Note this is not the answer for wasting mem, but it does improve some corner cases.
+    */
+    size_t min_primary_size = 0;
+    if (4 * chunk_size >= s_primary_min_cutoff) {
+        min_primary_size = s_primary_min_cutoff;
+    } else {
+        min_primary_size = 0;
+    }
+
     size_t page_size = aws_system_info_page_size();
     /* TODO: if people override their allocator, this will not wrap their allocator at all. eg: mem-tracing allocator
      * will be ignored. */
@@ -248,6 +273,7 @@ struct aws_s3_buffer_pool *aws_s3_default_buffer_pool_new(
      * Tries to balance between how many allocations use buffer and buffer space
      * being wasted. */
     buffer_pool->primary_size_cutoff = chunk_size * 4;
+    buffer_pool->primary_size_min_cutoff = min_primary_size;
     buffer_pool->mem_limit = adjusted_mem_lim;
 
     int mutex_error = aws_mutex_init(&buffer_pool->mutex);
@@ -395,16 +421,18 @@ static void s_aws_ticket_wrapper_destroy(void *data) {
 
         if (ticket->ptr == NULL) {
             /* Ticket was never used, make sure to clean up reserved count. */
-            if (ticket->is_special_block) {
+            if (ticket->reserved_from == AWS_S3_BUFFER_POOL_RESERVED_FROM_SPECIAL) {
                 buffer_pool->special_blocks_reserved -= ticket->size;
-            } else if (ticket->size <= buffer_pool->primary_size_cutoff) {
+            } else if (ticket->reserved_from == AWS_S3_BUFFER_POOL_RESERVED_FROM_PRIMARY) {
                 buffer_pool->primary_reserved -= ticket->size;
-            } else {
+            } else if (ticket->reserved_from == AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY) {
                 buffer_pool->secondary_reserved -= ticket->size;
+            } else {
+                AWS_ASSERT(false); /* someone forgot to update this after adding new member. */
             }
         } else {
             /* Handle special block returns */
-            if (ticket->is_special_block) {
+            if (ticket->reserved_from == AWS_S3_BUFFER_POOL_RESERVED_FROM_SPECIAL) {
                 struct aws_hash_element *elem = NULL;
                 aws_hash_table_find(&buffer_pool->special_blocks, (void *)ticket->size, &elem);
                 /* TODO: the lifetime between the ticket and the special block, let's assume the block will always be
@@ -427,7 +455,7 @@ static void s_aws_ticket_wrapper_destroy(void *data) {
                 /* Make sure we do find where the ticket buffer from the list. */
                 AWS_FATAL_ASSERT(found);
                 buffer_pool->special_blocks_used -= ticket->size;
-            } else if (ticket->size <= buffer_pool->primary_size_cutoff) {
+            } else if (ticket->reserved_from == AWS_S3_BUFFER_POOL_RESERVED_FROM_PRIMARY) {
 
                 size_t chunks_used = ticket->size / buffer_pool->chunk_size;
                 if (ticket->size % buffer_pool->chunk_size != 0) {
@@ -455,13 +483,15 @@ static void s_aws_ticket_wrapper_destroy(void *data) {
                 if (ticket->forced) {
                     buffer_pool->forced_used -= ticket->size;
                 }
-            } else {
+            } else if (ticket->reserved_from == AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY) {
                 aws_mem_release(buffer_pool->base_allocator, ticket->ptr);
                 buffer_pool->secondary_used -= ticket->size;
 
                 if (ticket->forced) {
                     buffer_pool->forced_used -= ticket->size;
                 }
+            } else {
+                AWS_ASSERT(false); /* someone forgot to update this after adding new member. */
             }
         }
         aws_mem_release(buffer_pool->base_allocator, ticket);
@@ -609,10 +639,18 @@ struct aws_s3_default_buffer_ticket *s_try_reserve_synced(
         /* Check if this is a special-sized allocation */
         if (from_special) {
             ticket->is_special_block = true;
+            ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_SPECIAL;
             buffer_pool->special_blocks_reserved += meta.size;
-        } else if (meta.size <= buffer_pool->primary_size_cutoff) {
-            buffer_pool->primary_reserved += meta.size;
+        } else if (meta.size <= buffer_pool->primary_size_cutoff) { /* TODO: this needs to be smarter. what if block size more than mem_lim*/
+            if (meta.size <= buffer_pool->primary_size_min_cutoff) {
+                ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY;
+                buffer_pool->secondary_reserved += meta.size;
+            } else {
+                ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_PRIMARY;
+                buffer_pool->primary_reserved += meta.size;
+            }
         } else {
+            ticket->reserved_from = AWS_S3_BUFFER_POOL_RESERVED_FROM_SECONDARY;
             buffer_pool->secondary_reserved += meta.size;
         }
     }

tests/CMakeLists.txt

@@ -447,6 +447,7 @@ add_test_case(test_s3_buffer_pool_special_size_reuse)
 add_test_case(test_s3_buffer_pool_special_size_with_limits)
 add_test_case(test_s3_buffer_pool_special_size_mixed)
 add_test_case(test_s3_buffer_pool_special_size_trim)
+add_test_case(test_s3_buffer_pool_reserve_tiny_chunks)
 
 add_net_test_case(client_update_upload_part_timeout)
 add_net_test_case(client_meta_request_override_part_size)

tests/s3_default_buffer_pool_tests.c

@@ -635,3 +635,38 @@ static int s_test_s3_buffer_pool_forced_buffer_wont_stop_reservations(struct aws
 AWS_TEST_CASE(
     test_s3_buffer_pool_forced_buffer_wont_stop_reservations,
     s_test_s3_buffer_pool_forced_buffer_wont_stop_reservations)
+
+
+/* Test small requests reserve expected amount of mem without rounding up to part size. */
+static int s_test_s3_buffer_pool_reserve_tiny_chunks(struct aws_allocator *allocator, void *ctx) {
+    (void)ctx;
+    const size_t chunk_size = MB_TO_BYTES(8);
+    const size_t small_size = MB_TO_BYTES(1);
+    const size_t mem_limit = GB_TO_BYTES(1);
+    struct aws_s3_buffer_pool *buffer_pool = aws_s3_default_buffer_pool_new(
+        allocator, (struct aws_s3_buffer_pool_config){.part_size = chunk_size, .memory_limit = mem_limit});
+
+    struct aws_s3_buffer_ticket *ticket = NULL;
+    struct aws_future_s3_buffer_ticket *future =
+        aws_s3_default_buffer_pool_reserve(buffer_pool, (struct aws_s3_buffer_pool_reserve_meta){.size = small_size});
+    ASSERT_TRUE(aws_future_s3_buffer_ticket_is_done(future));
+    ASSERT_INT_EQUALS(aws_future_s3_buffer_ticket_get_error(future), AWS_OP_SUCCESS);
+    normal_ticket = aws_future_s3_buffer_ticket_get_result_by_move(future);
+    struct aws_byte_buf buf = aws_s3_buffer_ticket_claim(ticket);
+    ASSERT_UINT_EQUALS(small, buf.capacity);
+
+    struct aws_s3_default_buffer_pool_usage_stats stats = aws_s3_default_buffer_pool_get_usage(buffer_pool);
+    ASSERT_INT_EQUALS(small_size, stats.secondary_used);
+
+    aws_s3_buffer_ticket_release(ticket);
+    aws_future_s3_buffer_ticket_release(future);
+
+    ASSERT_INT_EQUALS(0, stats.secondary_used);
+
+    /* Cleanup */
+    aws_s3_default_buffer_pool_destroy(buffer_pool);
+    return 0;
+}
+AWS_TEST_CASE(
+    test_s3_buffer_pool_reserve_tiny_chunks,
+    s_test_s3_buffer_pool_reserve_tiny_chunks)