From 931d523dccc7f276e0a665efd7168ff2c72b579a Mon Sep 17 00:00:00 2001
From: daanx <daan@effp.org>
Date: Thu, 29 Feb 2024 18:17:58 -0800
Subject: [PATCH] update mstress to let the main thread participate

---
 src/alloc.c        |  5 ++++-
 src/segment.c      |  4 ++--
 test/test-stress.c | 12 +++++++-----
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/alloc.c b/src/alloc.c
index 7135a857..e2273d28 100644
--- a/src/alloc.c
+++ b/src/alloc.c
@@ -408,7 +408,10 @@ static mi_decl_noinline void _mi_free_block_mt(mi_page_t* page, mi_block_t* bloc
 {
   // first see if the segment was abandoned and we can reclaim it
   mi_segment_t* const segment = _mi_page_segment(page);
-  if (mi_option_is_enabled(mi_option_abandoned_reclaim_on_free) &&
+  if (mi_option_is_enabled(mi_option_abandoned_reclaim_on_free) && 
+      #if MI_HUGE_PAGE_ABANDON
+      segment->page_kind != MI_PAGE_HUGE && 
+      #endif
       mi_atomic_load_relaxed(&segment->thread_id) == 0) 
   {
     // the segment is abandoned, try to reclaim it into our heap
diff --git a/src/segment.c b/src/segment.c
index 3f325dcf..3dd37429 100644
--- a/src/segment.c
+++ b/src/segment.c
@@ -867,7 +867,6 @@ static mi_segment_t* mi_segment_reclaim(mi_segment_t* segment, mi_heap_t* heap,
       // set the heap again and allow heap thread delayed free again.
       mi_page_set_heap(page, heap);
       _mi_page_use_delayed_free(page, MI_USE_DELAYED_FREE, true); // override never (after heap is set)
-      // TODO: should we not collect again given that we just collected in `check_free`?
       _mi_page_free_collect(page, false); // ensure used count is up to date
       if (mi_page_all_free(page)) {
         // if everything free already, clear the page directly
@@ -906,8 +905,9 @@ static mi_segment_t* mi_segment_reclaim(mi_segment_t* segment, mi_heap_t* heap,
 bool _mi_segment_attempt_reclaim(mi_heap_t* heap, mi_segment_t* segment) {
   if (mi_atomic_load_relaxed(&segment->thread_id) != 0) return false;  // it is not abandoned  
   if (_mi_arena_segment_clear_abandoned(segment->memid)) {  // atomically unabandon
+    mi_atomic_decrement_relaxed(&abandoned_count);
     mi_segment_t* res = mi_segment_reclaim(segment, heap, 0, NULL, &heap->tld->segments);
-    mi_assert_internal(res != NULL);
+    mi_assert_internal(res == segment);
     return (res != NULL);
   }
   return false;
diff --git a/test/test-stress.c b/test/test-stress.c
index 3ecb67bd..8c81c009 100644
--- a/test/test-stress.c
+++ b/test/test-stress.c
@@ -291,13 +291,14 @@ static void run_os_threads(size_t nthreads, void (*fun)(intptr_t)) {
   thread_entry_fun = fun;
   DWORD* tids = (DWORD*)custom_calloc(nthreads,sizeof(DWORD));
   HANDLE* thandles = (HANDLE*)custom_calloc(nthreads,sizeof(HANDLE));
-  for (uintptr_t i = 0; i < nthreads; i++) {
+  for (uintptr_t i = 1; i < nthreads; i++) {
     thandles[i] = CreateThread(0, 8*1024, &thread_entry, (void*)(i), 0, &tids[i]);
   }
-  for (size_t i = 0; i < nthreads; i++) {
+  fun(0); // run the main thread as well
+  for (size_t i = 1; i < nthreads; i++) {
     WaitForSingleObject(thandles[i], INFINITE);
   }
-  for (size_t i = 0; i < nthreads; i++) {
+  for (size_t i = 1; i < nthreads; i++) {
     CloseHandle(thandles[i]);
   }
   custom_free(tids);
@@ -325,10 +326,11 @@ static void run_os_threads(size_t nthreads, void (*fun)(intptr_t)) {
   pthread_t* threads = (pthread_t*)custom_calloc(nthreads,sizeof(pthread_t));
   memset(threads, 0, sizeof(pthread_t) * nthreads);
   //pthread_setconcurrency(nthreads);
-  for (size_t i = 0; i < nthreads; i++) {
+  for (size_t i = 1; i < nthreads; i++) {
     pthread_create(&threads[i], NULL, &thread_entry, (void*)i);
   }
-  for (size_t i = 0; i < nthreads; i++) {
+  fun(0); // run the main thread as well
+  for (size_t i = 1; i < nthreads; i++) {
     pthread_join(threads[i], NULL);
   }
   custom_free(threads);