From 12257e5bc1f3c201e870000e501552475ae33e69 Mon Sep 17 00:00:00 2001
From: Daan Leijen <Daan@microsoft.com>
Date: Wed, 21 Aug 2019 08:06:40 -0700
Subject: [PATCH] fix thread-id reset on page abandonment

---
 include/mimalloc-internal.h |  8 +++++++-
 src/init.c                  |  5 ++++-
 src/segment.c               | 10 ++++++++--
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/include/mimalloc-internal.h b/include/mimalloc-internal.h
index e4bb1387..6455d57e 100644
--- a/include/mimalloc-internal.h
+++ b/include/mimalloc-internal.h
@@ -323,7 +323,13 @@ static inline uintptr_t mi_page_thread_id(const mi_page_t* page) {
 }
 
 static inline void mi_page_init_flags(mi_page_t* page, uintptr_t thread_id) {
-  page->flags = thread_id;  
+  mi_assert_internal((thread_id & MI_PAGE_FLAGS_MASK) == 0);
+  page->flags = thread_id;
+}
+
+static inline void mi_page_set_thread_id(mi_page_t* page, uintptr_t thread_id) {
+  mi_assert_internal((thread_id & MI_PAGE_FLAGS_MASK) == 0);
+  page->flags = thread_id | (page->flags & MI_PAGE_FLAGS_MASK);
 }
 
 static inline bool mi_page_is_in_full(const mi_page_t* page) {
diff --git a/src/init.c b/src/init.c
index 13aba9b3..37afbe72 100644
--- a/src/init.c
+++ b/src/init.c
@@ -231,6 +231,8 @@ static bool _mi_heap_done(void) {
   mi_heap_t* heap = _mi_heap_default;
   if (!mi_heap_is_initialized(heap)) return true;
 
+  void* p = mi_malloc(16);
+
   // reset default heap
   _mi_heap_default = (_mi_is_main_thread() ? &_mi_heap_main : (mi_heap_t*)&_mi_heap_empty);
 
@@ -239,7 +241,7 @@ static bool _mi_heap_done(void) {
   // switch to backing heap and free it
   heap = heap->tld->heap_backing;
   if (!mi_heap_is_initialized(heap)) return false;
-
+  
   // collect if not the main thread
   if (heap != &_mi_heap_main) {
     _mi_heap_collect_abandon(heap);
@@ -258,6 +260,7 @@ static bool _mi_heap_done(void) {
     mi_assert_internal(heap->tld->heap_backing == &_mi_heap_main);
   }
 #endif
+  mi_free(p);
   return false;
 }
 
diff --git a/src/segment.c b/src/segment.c
index 78d5d0bd..d5a2288a 100644
--- a/src/segment.c
+++ b/src/segment.c
@@ -549,6 +549,11 @@ static void mi_segment_abandon(mi_segment_t* segment, mi_segments_tld_t* tld) {
   mi_assert_internal(segment->used > 0);
   mi_assert_internal(segment->abandoned_next == NULL);
   mi_assert_expensive(mi_segment_is_valid(segment));
+#if MI_DEBUG>1
+  for (size_t i = 0; i < segment->capacity; i++) {
+    mi_assert_internal(!segment->pages[i].segment_in_use || mi_page_thread_id(&segment->pages[i]) == 0);
+  }
+#endif
   // remove the segment from the free page queue if needed
   mi_segment_remove_from_free_queue(segment,tld);
   mi_assert_internal(segment->next == NULL && segment->prev == NULL);
@@ -563,10 +568,11 @@ static void mi_segment_abandon(mi_segment_t* segment, mi_segments_tld_t* tld) {
 }
 
 void _mi_segment_page_abandon(mi_page_t* page, mi_segments_tld_t* tld) {
-  mi_assert(page != NULL);
+  mi_assert(page != NULL && mi_page_thread_id(page) != 0);
   mi_segment_t* segment = _mi_page_segment(page);
   mi_assert_expensive(mi_segment_is_valid(segment));
   segment->abandoned++;
+  mi_page_set_thread_id(page, 0);
   _mi_stat_increase(&tld->stats->pages_abandoned, 1);
   mi_assert_internal(segment->abandoned <= segment->used);
   if (segment->used == segment->abandoned) {
@@ -618,7 +624,7 @@ bool _mi_segment_try_reclaim_abandoned( mi_heap_t* heap, bool try_all, mi_segmen
         }
         else {
           // otherwise reclaim it
-          mi_page_init_flags(page,segment->thread_id);
+          mi_page_set_thread_id(page,segment->thread_id);
           _mi_page_reclaim(heap,page);
         }
       }