fix the try-trim logic (#621)

TingDaoK · web-flow · commit 7397f2fc077c · 2026-03-31T14:33:35.000-07:00
diff --git a/source/s3_default_buffer_pool.c b/source/s3_default_buffer_pool.c
@@ -434,26 +434,27 @@ static void s_aws_ticket_wrapper_destroy(void *data) {
             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
-                 * available first. */
-                AWS_FATAL_ASSERT(elem != NULL);
-                special_list = elem->value;
-
-                bool found = false;
-                for (size_t i = 0; i < aws_array_list_length(&special_list->blocks); ++i) {
-                    struct s3_buffer_pool_block *block;
-                    aws_array_list_get_at_ptr(&special_list->blocks, (void **)&block, i);
-                    if (block->block_ptr == ticket->ptr) {
-                        /* Make sure the block is marked as used. */
-                        AWS_FATAL_ASSERT(block->alloc_bit_mask == UINT16_MAX);
-                        block->alloc_bit_mask = 0;
-                        found = true;
-                        break;
+                if (elem) {
+                    special_list = elem->value;
+
+                    bool found = false;
+                    for (size_t i = 0; i < aws_array_list_length(&special_list->blocks); ++i) {
+                        struct s3_buffer_pool_block *block;
+                        aws_array_list_get_at_ptr(&special_list->blocks, (void **)&block, i);
+                        if (block->block_ptr == ticket->ptr) {
+                            /* Make sure the block is marked as used. */
+                            if (block->alloc_bit_mask == UINT16_MAX) {
+                                block->alloc_bit_mask = 0;
+                                found = true;
+                                break;
+                            }
+                        }
+                    }
+                    /* Make sure we do find where the ticket buffer from the list. */
+                    if (found) {
+                        buffer_pool->special_blocks_used -= ticket->size;
                     }
                 }
-                /* 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->reserved_from == AWS_S3_BUFFER_POOL_RESERVED_FROM_PRIMARY) {
 
                 size_t chunks_used = ticket->size / buffer_pool->chunk_size;
@@ -557,37 +558,31 @@ static void s_aws_ticket_wrapper_destroy(void *data) {
 static bool s_should_trim_for_reserve_synced(
     bool from_special,
     size_t to_reserve,
-    struct aws_s3_default_buffer_pool *buffer_pool) {
+    struct aws_s3_default_buffer_pool *buffer_pool,
+    size_t overall_taken) {
     if (from_special || to_reserve <= buffer_pool->primary_size_cutoff) {
         /* Only trim when it will be allocated from secondary, which is directly from malloc. */
         return false;
     }
-    /* The tracked usage SHOULD NOT overflow. */
-    size_t total_allocated =
-        buffer_pool->primary_allocated + buffer_pool->secondary_used + buffer_pool->special_blocks_allocated;
+
     size_t total_allocation_needs = 0;
-    if (aws_add_size_checked(total_allocated, to_reserve, &total_allocation_needs)) {
-        /* Will overflow, trim it. */
-        return true;
+    if (aws_add_size_checked(overall_taken, to_reserve, &total_allocation_needs)) {
+        /* Will overflow, ignores it, which should not happen. */
+        return false;
     }
     if (total_allocation_needs < buffer_pool->mem_limit) {
         /* No need to trim as we still have space to allocate the new block. */
         return false;
     }
 
-    size_t primary_overallocation = 0;
-    int overflow = 0;
-    overflow |=
-        aws_sub_size_checked(buffer_pool->primary_allocated, buffer_pool->primary_used, &primary_overallocation);
-    overflow |= aws_sub_size_checked(primary_overallocation, buffer_pool->primary_reserved, &primary_overallocation);
-    size_t special_overallocation = 0;
-    overflow |= aws_sub_size_checked(
-        buffer_pool->special_blocks_allocated, buffer_pool->special_blocks_used, &special_overallocation);
-    overflow |=
-        aws_sub_size_checked(special_overallocation, buffer_pool->special_blocks_reserved, &special_overallocation);
-    /* The already allocated should be reasonable and not cause overflow. Otherwise, bugs in the code. */
-    AWS_FATAL_ASSERT(!overflow);
-    /* Use max if overflow */
+    size_t primary_overallocation = aws_sub_size_saturating(buffer_pool->primary_allocated, buffer_pool->primary_used);
+    /* Reserved can be more tha allocated */
+    primary_overallocation = aws_sub_size_saturating(primary_overallocation, buffer_pool->primary_reserved);
+    size_t special_overallocation =
+        aws_sub_size_saturating(buffer_pool->special_blocks_allocated, buffer_pool->special_blocks_used);
+    /* Reserved can be more tha allocated */
+    special_overallocation = aws_sub_size_saturating(special_overallocation, buffer_pool->special_blocks_reserved);
+
     size_t total_overallocation = aws_add_size_saturating(special_overallocation, primary_overallocation);
     if (total_overallocation < to_reserve) {
         /* If the overallocation is less than the new block, trim it won't help, skip trimming. */
@@ -616,7 +611,7 @@ struct aws_s3_default_buffer_ticket *s_try_reserve_synced(
      * the blocks, trim the blocks in hopes we can free up enough memory.
      * TODO: something smarter, like partial trim?
      */
-    if (s_should_trim_for_reserve_synced(from_special, meta.size, buffer_pool)) {
+    if (s_should_trim_for_reserve_synced(from_special, meta.size, buffer_pool, overall_taken)) {
         s_buffer_pool_trim_synced(buffer_pool);
         overall_taken = buffer_pool->primary_used + buffer_pool->primary_reserved + buffer_pool->secondary_used +
                         buffer_pool->secondary_reserved + buffer_pool->special_blocks_reserved +
@@ -627,7 +622,7 @@ struct aws_s3_default_buffer_ticket *s_try_reserve_synced(
     const size_t max_impact_of_forced_on_limit =
         (size_t)(buffer_pool->mem_limit * (s_max_impact_of_forced_buffers_on_memory_limit_as_percentage / 100.0));
     if (buffer_pool->forced_used > max_impact_of_forced_on_limit) {
-        overall_taken -= buffer_pool->forced_used - max_impact_of_forced_on_limit;
+        overall_taken = (overall_taken - buffer_pool->forced_used) + max_impact_of_forced_on_limit;
     }
 
     if ((meta.size + overall_taken) <= buffer_pool->mem_limit) {
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -448,6 +448,7 @@ 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_test_case(test_s3_buffer_pool_trim_reserved_but_unallocated)
 
 add_net_test_case(client_update_upload_part_timeout)
 add_net_test_case(client_meta_request_override_part_size)
diff --git a/tests/s3_default_buffer_pool_tests.c b/tests/s3_default_buffer_pool_tests.c
@@ -679,3 +679,72 @@ static int s_test_s3_buffer_pool_reserve_tiny_chunks(struct aws_allocator *alloc
     return 0;
 }
 AWS_TEST_CASE(test_s3_buffer_pool_reserve_tiny_chunks, s_test_s3_buffer_pool_reserve_tiny_chunks)
+
+/*
+ * Regression test for s_should_trim_for_reserve_synced when buffers are reserved but not yet claimed.
+ *
+ * When tickets are reserved but never claimed, primary_allocated stays 0 while primary_reserved > 0.
+ * The overallocation calculation (primary_allocated - primary_used - primary_reserved) must not
+ * underflow -- it should saturate to 0. This also verifies that a subsequent secondary reserve
+ * still succeeds (i.e. the function correctly returns false / no crash).
+ */
+static int s_test_s3_buffer_pool_trim_reserved_but_unallocated(struct aws_allocator *allocator, void *ctx) {
+    (void)ctx;
+    const size_t chunk_size = MB_TO_BYTES(8);
+    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,
+        });
+    ASSERT_NOT_NULL(buffer_pool);
+
+    /* Reserve primary-sized tickets WITHOUT claiming them.
+     * This leaves primary_reserved > 0 but primary_allocated == 0. */
+    const size_t num_reserved = 50;
+    struct aws_future_s3_buffer_ticket *reserved_futures[50];
+    struct aws_s3_buffer_ticket *reserved_tickets[50];
+    for (size_t i = 0; i < num_reserved; ++i) {
+        reserved_futures[i] = aws_s3_default_buffer_pool_reserve(
+            buffer_pool, (struct aws_s3_buffer_pool_reserve_meta){.size = chunk_size});
+        ASSERT_TRUE(aws_future_s3_buffer_ticket_is_done(reserved_futures[i]));
+        ASSERT_INT_EQUALS(aws_future_s3_buffer_ticket_get_error(reserved_futures[i]), AWS_OP_SUCCESS);
+        reserved_tickets[i] = aws_future_s3_buffer_ticket_get_result_by_move(reserved_futures[i]);
+        ASSERT_NOT_NULL(reserved_tickets[i]);
+        /* Intentionally do NOT call aws_s3_buffer_ticket_claim() -- ticket stays in reserve stage. */
+    }
+
+    /* Sanity check: reserved > 0, allocated == 0 */
+    struct aws_s3_default_buffer_pool_usage_stats stats = aws_s3_default_buffer_pool_get_usage(buffer_pool);
+    ASSERT_TRUE(stats.primary_reserved > 0);
+    ASSERT_UINT_EQUALS(0, stats.primary_allocated);
+
+    /* Now reserve a large secondary buffer. This triggers s_should_trim_for_reserve_synced.
+     * The overallocation calculation must saturate to 0 (not crash/underflow).
+     * The reserve must succeed since mem_limit is not exhausted. */
+    struct aws_future_s3_buffer_ticket *secondary_future = aws_s3_default_buffer_pool_reserve(
+        buffer_pool,
+        (struct aws_s3_buffer_pool_reserve_meta){
+            .size = MB_TO_BYTES(800),
+        });
+    ASSERT_NOT_NULL(secondary_future);
+    /* Release all the reserved ones. */
+    for (size_t i = 0; i < num_reserved; ++i) {
+        aws_s3_buffer_ticket_release(reserved_tickets[i]);
+        aws_future_s3_buffer_ticket_release(reserved_futures[i]);
+    }
+    ASSERT_TRUE(aws_future_s3_buffer_ticket_is_done(secondary_future));
+    ASSERT_INT_EQUALS(aws_future_s3_buffer_ticket_get_error(secondary_future), AWS_OP_SUCCESS);
+    struct aws_s3_buffer_ticket *secondary_ticket = aws_future_s3_buffer_ticket_get_result_by_move(secondary_future);
+    ASSERT_NOT_NULL(secondary_ticket);
+    struct aws_byte_buf buf = aws_s3_buffer_ticket_claim(secondary_ticket);
+    ASSERT_NOT_NULL(buf.buffer);
+
+    /* Cleanup */
+    aws_s3_buffer_ticket_release(secondary_ticket);
+    aws_future_s3_buffer_ticket_release(secondary_future);
+    aws_s3_default_buffer_pool_destroy(buffer_pool);
+    return 0;
+}
+AWS_TEST_CASE(test_s3_buffer_pool_trim_reserved_but_unallocated, s_test_s3_buffer_pool_trim_reserved_but_unallocated)
