From a7bd9c08c8ad7862d138c97f795a7b607b1f4c21 Mon Sep 17 00:00:00 2001
From: Daan Leijen <daan@effp.org>
Date: Wed, 23 Nov 2022 09:58:45 -0800
Subject: [PATCH] fix decommit of huge pages

---
 src/segment.c          | 35 +++++++++++++++++------------------
 test/main-override.cpp | 22 ++++++++++++++++++----
 2 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/src/segment.c b/src/segment.c
index be496f21..85cac395 100644
--- a/src/segment.c
+++ b/src/segment.c
@@ -1224,26 +1224,22 @@ static mi_page_t* mi_segment_huge_page_alloc(size_t size, size_t page_alignment,
   mi_page_t* page = mi_segment_find_free(segment, tld);
   mi_assert_internal(page != NULL);
 
-  if (page_alignment > 0) {
-    size_t psize;
-    size_t pre_size;
-    uint8_t* p = (uint8_t*)_mi_segment_page_start(segment, page, 0, &psize, &pre_size);
-    uint8_t* aligned_p = (uint8_t*)_mi_align_up((uintptr_t)p, page_alignment);
-    mi_assert_internal(_mi_is_aligned(aligned_p, page_alignment));
-    mi_assert_internal(psize - (aligned_p - p) >= size);
-    if (!segment->mem_is_pinned && page->is_committed) {
-      // decommit the part of the page that is unused; this can be quite large (close to MI_SEGMENT_SIZE)
-      uint8_t* decommit_start = p + sizeof(mi_block_t); // for the free list
-      ptrdiff_t decommit_size = aligned_p - decommit_start;
-      _mi_os_reset(decommit_start, decommit_size, os_tld->stats);
-    }
-  }
-
   // for huge pages we initialize the xblock_size as we may
   // overallocate to accommodate large alignments.
   size_t psize;
-  _mi_segment_page_start(segment, page, 0, &psize, NULL);
+  uint8_t* start = _mi_segment_page_start(segment, page, 0, &psize, NULL);
   page->xblock_size = (psize > MI_HUGE_BLOCK_SIZE ? MI_HUGE_BLOCK_SIZE : (uint32_t)psize);
+
+  // reset the part of the page that will not be used; this can be quite large (close to MI_SEGMENT_SIZE)
+  if (page_alignment > 0 && !segment->mem_is_pinned && page->is_committed) {
+    uint8_t* aligned_p = (uint8_t*)_mi_align_up((uintptr_t)start, page_alignment);
+    mi_assert_internal(_mi_is_aligned(aligned_p, page_alignment));
+    mi_assert_internal(psize - (aligned_p - start) >= size);
+    uint8_t* decommit_start = start + sizeof(mi_block_t); // for the free list
+    ptrdiff_t decommit_size = aligned_p - decommit_start;
+    _mi_os_reset(decommit_start, decommit_size, os_tld->stats);  // do not decommit as it may be in a region
+  }
+
   return page;
 }
 
@@ -1283,8 +1279,11 @@ void _mi_segment_huge_page_reset(mi_segment_t* segment, mi_page_t* page, mi_bloc
   mi_assert_internal(segment == _mi_page_segment(page));
   mi_assert_internal(page->used == 1); // this is called just before the free
   mi_assert_internal(page->free == NULL);
-  const size_t bsize = mi_page_block_size(page);
-  _mi_os_reset(block + 1, bsize - sizeof(mi_block_t), &_mi_stats_main);
+  if (!segment->mem_is_pinned && page->is_committed) {
+    const size_t usize = mi_usable_size(block) - sizeof(mi_block_t);
+    uint8_t* p = (uint8_t*)block + sizeof(mi_block_t);
+    _mi_os_reset(p, usize, &_mi_stats_main); 
+  }
 }
 #endif
 
diff --git a/test/main-override.cpp b/test/main-override.cpp
index f5cb3668..81f57298 100644
--- a/test/main-override.cpp
+++ b/test/main-override.cpp
@@ -34,23 +34,26 @@ static void various_tests();
 static void test_mt_shutdown();
 static void fail_aslr();              // issue #372
 static void tsan_numa_test();         // issue #414
-static void strdup_test();     // issue #445
+static void strdup_test();            // issue #445
+static void heap_thread_free_huge();
 
 static void test_stl_allocators();
 
 int main() {
   mi_stats_reset();  // ignore earlier allocations
+
+  heap_thread_free_huge();
+  /*
   heap_thread_free_large();
   heap_no_delete();
   heap_late_free();
   padding_shrink();
   various_tests();
   tsan_numa_test();
-  strdup_test();
-
+  strdup_test();  
   test_stl_allocators();
-
   test_mt_shutdown();
+  */
   //fail_aslr();
   mi_stats_print(NULL);
   return 0;
@@ -235,6 +238,17 @@ static void heap_thread_free_large() {
   }
 }
 
+static void heap_thread_free_huge_worker() {
+  mi_free(shared_p);
+}
+
+static void heap_thread_free_huge() {
+  for (int i = 0; i < 10; i++) {
+    shared_p = mi_malloc(1024 * 1024 * 1024);
+    auto t1 = std::thread(heap_thread_free_large_worker);
+    t1.join();
+  }
+}
 
 
 static void test_mt_shutdown()