From df48b8b6c9c53ed1fae973e5e941f7a5f45e39c7 Mon Sep 17 00:00:00 2001
From: Sergiy Kuryata <sergeyk@microsoft.com>
Date: Tue, 17 Dec 2024 13:17:28 -0800
Subject: [PATCH] Ensure large and medium segments are droped when allocating a
 new segment for medium and large objects

---
 src/heap.c | 44 +++++++++++++++++++++++---------------------
 1 file changed, 23 insertions(+), 21 deletions(-)

diff --git a/src/heap.c b/src/heap.c
index d84dcbbe..13722de6 100644
--- a/src/heap.c
+++ b/src/heap.c
@@ -656,7 +656,7 @@ bool mi_heap_visit_blocks(const mi_heap_t* heap, bool visit_blocks, mi_block_vis
 mi_segment_t* mi_segments_get_segment_to_drop_by_slice(mi_segments_tld_t* tld, size_t alloc_block_size);
 const mi_slice_t* mi_segment_slices_end(const mi_segment_t* segment);
 
-static mi_segment_t* mi_heap_get_segment_to_drop(mi_heap_t* heap, size_t alloc_block_size) {
+static mi_segment_t* mi_heap_get_segment_to_drop(mi_heap_t* heap, size_t targetSegmentCount, size_t alloc_block_size) {
     mi_segment_t* segment = NULL;
 
     mi_page_kind_t page_kind = mi_page_kind_from_size(alloc_block_size);
@@ -667,26 +667,28 @@ static mi_segment_t* mi_heap_get_segment_to_drop(mi_heap_t* heap, size_t alloc_b
         return heap->tld->segments.medium_segment;
     }
 
-    int i = 0;
-    mi_page_queue_t* fullPageQueue = &heap->pages[MI_BIN_FULL];
-    for (mi_page_t* page = fullPageQueue->first; page != NULL; page = page->next) {
-        mi_segment_t* temp_segment = _mi_ptr_segment(page);
-        if (temp_segment->page_kind == MI_PAGE_SMALL) {
-            if (segment == NULL) {
-                segment = temp_segment;
+    if (heap->tld->segments.count >= targetSegmentCount) {
+        int i = 0;
+        mi_page_queue_t* fullPageQueue = &heap->pages[MI_BIN_FULL];
+        for (mi_page_t* page = fullPageQueue->first; page != NULL; page = page->next) {
+            mi_segment_t* temp_segment = _mi_ptr_segment(page);
+            if (temp_segment->page_kind == MI_PAGE_SMALL) {
+                if (segment == NULL) {
+                    segment = temp_segment;
+                }
+                else if (temp_segment->used > segment->used) {
+                    segment = temp_segment;
+                }
+                if (i > 3) {
+                    break;
+                }
+                i++;
             }
-            else if (temp_segment->used > segment->used) {
-                segment = temp_segment;
-            }
-            if (i > 3) {
-                break;
-            }
-            i++;
         }
-    }
 
-    if (segment == NULL) {
-        segment = mi_segments_get_segment_to_drop_by_slice(&heap->tld->segments, alloc_block_size);
+        if (segment == NULL) {
+            segment = mi_segments_get_segment_to_drop_by_slice(&heap->tld->segments, alloc_block_size);
+        }
     }
 
     return segment;
@@ -733,7 +735,7 @@ void mi_heap_drop_segment(mi_heap_t* heap, size_t targetSegmentCount, size_t all
     while (heap->tld->segments.count >= targetSegmentCount) {
 
         // 1. Find a segment to drop (abandon) using the Full Page queue
-        mi_segment_t* segmentToAbandon = mi_heap_get_segment_to_drop(heap, alloc_block_size);
+        mi_segment_t* segmentToAbandon = mi_heap_get_segment_to_drop(heap, targetSegmentCount, alloc_block_size);
         if (segmentToAbandon == NULL) {
             break;
         }
@@ -769,8 +771,8 @@ void mi_heap_drop_segment_if_required(mi_heap_t* heap, size_t alloc_block_size)
 {
     size_t targetSegmentCount = mi_option_get_size(mi_option_max_segments_per_heap);
     if ((targetSegmentCount > 0) && 
-        (alloc_block_size <= MI_LARGE_OBJ_SIZE_MAX) &&
-        (heap->tld->segments.count >= targetSegmentCount)) {
+        ((alloc_block_size > MI_SMALL_OBJ_SIZE_MAX) ||
+        (heap->tld->segments.count >= targetSegmentCount))) {
 
         mi_heap_drop_segment(heap, targetSegmentCount, alloc_block_size);
     }