profile: sample the small allocation that refilled its page, new threads, and the blocks mi_heap_destroy frees

CMakeLists.txt

@@ -1173,7 +1173,9 @@ if (MI_BUILD_TESTS)
     set(mi_test_env "")
     if(MI_GUARDED AND TEST_NAME STREQUAL "purge-holes")
       list(APPEND mi_test_env MIMALLOC_GUARDED_SAMPLE_RATE=0)     # it knows the page layout, which guard pages change (set explicitly: the CI sets a rate for the whole job)
-    elseif(MI_GUARDED AND NOT (TEST_NAME STREQUAL "profile"))
+    elseif(MI_GUARDED AND TEST_NAME STREQUAL "profile")
+      list(APPEND mi_test_env MIMALLOC_GUARDED_SAMPLE_RATE=0)     # it checks what the samples add up to, which guarded samples in between change (set explicitly: the CI sets a rate for the whole job); the case with guarded blocks sets its own rate
+    elseif(MI_GUARDED)
       list(APPEND mi_test_env MIMALLOC_GUARDED_SAMPLE_RATE=1024)
     elseif(TEST_NAME STREQUAL "stress-heaps")
       list(APPEND mi_test_env MIMALLOC_ARENA_EAGER_COMMIT=0 MIMALLOC_PAGE_COMMIT_ON_DEMAND=0)

include/mimalloc/internal.h

@@ -408,6 +408,7 @@ size_t        _mi_theap_update_sample_rate(mi_theap_t* theap);
 
 mi_decl_restrict void* _mi_theap_malloc_profiled(mi_theap_t* theap, size_t size, uint64_t requested_since_last_sample, bool zero, mi_page_t** ppage) mi_attr_noexcept;
 void          _mi_page_profile_on_free(mi_page_t* page, mi_block_t* block, void* p);
+void          _mi_page_profile_free_all(const mi_heap_area_t* area, mi_page_t* page);
 size_t        _mi_theap_set_profile_sample_rate(mi_theap_t* theap, size_t sample_rate);
 
 

src/arena.c

@@ -3060,6 +3060,9 @@ static bool mi_heap_delete_page(const mi_heap_t* heap, const mi_heap_area_t* are
     #if MI_GUARDED
     _mi_page_unguard_all(page);          // remove potential interior guard pages 
     #endif
+    #if MI_PROFILE
+    _mi_page_profile_free_all(area,page); // the sampled blocks that are still in use get their `on_free`
+    #endif
     // destroy the page
     // Drop what is (still) on the thread-free list first: `_mi_arenas_page_free` collects it, and with the used
     // count reset that reads as a corrupted list (more blocks freed than were in use).

src/page.c

@@ -1834,6 +1834,7 @@ static bool mi_page_extend_free(mi_theap_t* theap, mi_page_t* page) {
   //uint8_t* page_start =
   mi_page_area(page, &page_size);
   mi_theap_stat_counter_increase(theap, pages_extended, 1);
+  mi_page_update_sample_countdown(page);  // the blocks of the previous extension are handed out: count them
 
   // calculate the extend count
   const size_t bsize = mi_page_block_size(page);
@@ -2368,6 +2369,17 @@ void* _mi_malloc_generic(mi_theap_t* theap, size_t size, size_t zero_huge_alignm
         page = mi_page_queue_find_free(theap,pq);
         // mi_assert_internal(mi_page_block_size(page) <= MI_SMALL_MAX_OBJ_SIZE);
         if (page!=NULL) {        
+          #if MI_SAMPLE==1
+          // The blocks that the fast path took from a page are counted against the sample countdown when its free
+          // list is refilled, which is what `mi_page_queue_find_free` just did. If that used up the countdown, the
+          // allocation in progress (of the size class of those blocks) is the sample. If we leave it to the next 
+          // allocation that comes through here, that is one above `MI_SMALL_SIZE_MAX` (these always do) far more 
+          // often than a small one: with one such allocation for every 40 blocks of 64 bytes, less than 2% of the 
+          // sampled bytes went to the blocks of 64 bytes (`test-profile.c:test_profiler_small_attribution`).
+          if mi_unlikely(theap->sample_rate!=0 && mi_theap_should_sample(theap,req_size)) {
+            return _mi_theap_malloc_sampled(theap,req_size,zero,ppage);
+          }
+          #endif
           if (ppage!=NULL) { *ppage = page; }
           mi_assert_internal(mi_page_immediate_available(page)); // we should never recurse in _mi_page_malloc_zero
           return _mi_page_malloc_zero(theap,page,size,zero);

src/sample-profile.c

@@ -115,6 +115,34 @@ static mi_profiler_t* mi_theap_get_enabled_profiler(const mi_theap_t* theap) {
   }
 }
 
+//----------------------------------------------------------------------------
+// Layout of a profiled block: 
+//   [MI_BLOCK_TAG_PROFILED] [check] [mi_profiler_sample_data_t ... user data ...] [ ... block as the program sees it ... ]
+// `mi_free` recognizes it by the tag and by getting an interior pointer. `mi_heap_destroy` walks
+// the blocks of a page and has no pointer to go by: the first word of a regular block in use is
+// program data and can hold the value of the tag. The `check` word is the address of the block 
+// xor a random key, so a regular block is not mistaken for a profiled one.
+//-----------------------------------------------------------------------------
+
+static _Atomic(uintptr_t) mi_profile_check_key;  // = 0
+
+static uintptr_t mi_profile_block_check(mi_theap_t* theap, const mi_block_t* block) {
+  uintptr_t key = mi_atomic_load_relaxed(&mi_profile_check_key);
+  if mi_unlikely(key==0) {
+    if (theap==NULL) return 0;   // no block was profiled yet
+    uintptr_t expected = 0;
+    key = _mi_theap_random_next(theap) | 1;
+    if (!mi_atomic_cas_strong_acq_rel(&mi_profile_check_key, &expected, key)) { key = expected; }
+  }
+  return ((uintptr_t)block ^ key);
+}
+
+#define MI_PROFILE_SAMPLE_DATA_OFFSET  (sizeof(mi_block_t) + sizeof(uintptr_t))
+
+static inline size_t mi_profile_user_offset(size_t sample_user_data_size) {
+  return _mi_align_up(MI_PROFILE_SAMPLE_DATA_OFFSET + sizeof(mi_profiler_sample_data_t) + sample_user_data_size, MI_MAX_ALIGN_SIZE);
+}
+
 //----------------------------------------------------------------------------
 // Profile an allocation 
 //-----------------------------------------------------------------------------
@@ -142,11 +170,10 @@ mi_decl_noinline mi_decl_restrict void* _mi_theap_malloc_profiled(mi_theap_t* th
   }
   else {
     // overallocate a larger block to store the profiler data
-    // [MI_BLOCK_TAG_PROFILE] [usable size] [ ... profile data ... ] [... user data ...]
-    const size_t sample_data_offset    = sizeof(mi_block_t);
+    // [MI_BLOCK_TAG_PROFILED] [check] [ ... profile data ... ] [... user data ...]
+    const size_t sample_data_offset    = MI_PROFILE_SAMPLE_DATA_OFFSET;
     const size_t sample_user_data_size = _mi_align_up(prof->sample_data_size > MI_PROFILE_SAMPLE_DATA_MAX_SIZE ? MI_PROFILE_SAMPLE_DATA_MAX_SIZE : prof->sample_data_size, sizeof(void*)); 
-    const size_t sample_data_size      = sizeof(mi_profiler_sample_data_t) + sample_user_data_size;  // one void* too many just in case
-    const size_t user_offset           = _mi_align_up(sample_data_offset + sample_data_size, MI_MAX_ALIGN_SIZE);
+    const size_t user_offset           = mi_profile_user_offset(sample_user_data_size);
     const size_t oversize              = user_offset + size;
     mi_page_t* page = NULL;
     mi_block_t* const block = (mi_block_t*)_mi_malloc_generic_no_sample(theap,oversize,zero,&page); 
@@ -163,6 +190,7 @@ mi_decl_noinline mi_decl_restrict void* _mi_theap_malloc_profiled(mi_theap_t* th
     // Set up the profiled block as an interior pointer so the interior "slow" path is taken on mi_free (where we catch it to call on_free)
     mi_page_set_has_interior_pointers(page, true);
     block->next = MI_BLOCK_TAG_PROFILED;  
+    *((uintptr_t*)(block + 1)) = mi_profile_block_check(theap,block);
     p = (uint8_t*)block + user_offset;
     mi_profiler_sample_data_t* const sample_data = (mi_profiler_sample_data_t*)((uint8_t*)block + sample_data_offset);
     sample_data->user_data_size = sample_user_data_size;
@@ -174,23 +202,64 @@ mi_decl_noinline mi_decl_restrict void* _mi_theap_malloc_profiled(mi_theap_t* th
   }
   if (new_sample_rate!=0 && new_sample_rate != (size_t)theap->profile_sample_rate) { 
     _mi_theap_set_profile_sample_rate(theap,new_sample_rate);
+    // We are at a sample, where a period starts: start it with the new rate in full. (`_mi_theap_malloc_sampled`
+    // takes the period to be `sample_rate` long when it computes `bytes_since_last_sample`. A raised rate left the 
+    // countdown at the previous rate, and each sample reported the difference on top of the bytes that were requested.)
+    theap->profile_sample_countdown = theap->profile_sample_rate;
+    theap->sample_countdown = theap->sample_rate;
   }
   mi_theap_stat_counter_increase(theap,profile_samples,1);  
   return p;
 }
 
 void _mi_page_profile_on_free(mi_page_t* page, mi_block_t* block, void* p) {
+  MI_UNUSED_RELEASE(p);
   mi_assert_internal(mi_block_ptr_is_sampled(block,p));
+  uintptr_t* const pcheck = (uintptr_t*)(block + 1);
+  const uintptr_t check = mi_profile_block_check(NULL,block);
+  if mi_unlikely(check==0 || *pcheck != check) return;  // not a block that `_mi_theap_malloc_profiled` set up
+  // The free list only overwrites the tag. Clear the check word too, or the block could come back as a regular 
+  // block that starts with the value of the tag and still has the check behind it (see `_mi_page_profile_free_all`).
+  *pcheck = 0;
 
   // get the heap and profiler
   mi_heap_t* const heap = mi_page_heap(page);
   if (heap==NULL) return;
   mi_profiler_t* prof = mi_heap_profiler(heap);
   if (prof==NULL || !mi_profiler_is_enabled(prof) || prof->on_free==NULL) return;
 
-  // call the on_free callback
-  mi_profiler_sample_data_t* const sample_data = (mi_profiler_sample_data_t*)((uint8_t*)block + sizeof(mi_block_t));
-  (*prof->on_free)(prof, sample_data, p, heap);
+  // Call the on_free callback with the pointer that `on_alloc` got. `p` is another one if the block was allocated 
+  // with an alignment: `alloc-aligned.c` aligns inside the block that `_mi_theap_malloc_profiled` returned.
+  mi_profiler_sample_data_t* const sample_data = (mi_profiler_sample_data_t*)((uint8_t*)block + MI_PROFILE_SAMPLE_DATA_OFFSET);
+  void* const sampled_p = (uint8_t*)block + mi_profile_user_offset(sample_data->user_data_size);
+  mi_assert_internal((uint8_t*)p >= (uint8_t*)sampled_p);
+  (*prof->on_free)(prof, sample_data, sampled_p, heap);
+}
+
+// `mi_heap_destroy` frees the pages of a heap without a `mi_free` of each block: 
+// call `on_free` for the profiled blocks that are still in use in `page`.
+static bool mi_page_profile_free_visitor(const mi_heap_t* heap, const mi_heap_area_t* area, void* vblock, size_t block_size, void* arg) {
+  MI_UNUSED(heap); MI_UNUSED(area);
+  mi_page_t* const page = (mi_page_t*)arg;
+  mi_block_t* const block = (mi_block_t*)vblock;
+  if (block==NULL || block_size < MI_PROFILE_SAMPLE_DATA_OFFSET + sizeof(mi_profiler_sample_data_t)) return true;
+  if (block->next != MI_BLOCK_TAG_PROFILED) return true;
+  const uintptr_t check = mi_profile_block_check(NULL,block);
+  if (check==0 || *((const uintptr_t*)(block + 1)) != check) return true;
+  const mi_profiler_sample_data_t* const sample_data = (const mi_profiler_sample_data_t*)((uint8_t*)block + MI_PROFILE_SAMPLE_DATA_OFFSET);
+  const size_t user_offset = mi_profile_user_offset(sample_data->user_data_size);
+  if (user_offset >= block_size) return true;
+  _mi_page_profile_on_free(page, block, (uint8_t*)block + user_offset);  // clears the check word
+  return true;
+}
+
+void _mi_page_profile_free_all(const mi_heap_area_t* area, mi_page_t* page) {
+  if mi_likely(!mi_page_has_interior_pointers(page)) return;
+  mi_heap_t* const heap = mi_page_heap(page);
+  if (heap==NULL) return;
+  mi_profiler_t* const prof = mi_heap_profiler(heap);
+  if (prof==NULL || !mi_profiler_is_enabled(prof) || prof->on_free==NULL || prof->sample_data_size==0) return;
+  _mi_theap_area_visit_blocks(area, page, &mi_page_profile_free_visitor, page);
 }
 
 

src/theap.c

@@ -552,6 +552,21 @@ void _mi_theap_init(mi_theap_t* theap, mi_heap_t* heap, mi_tld_t* tld)
   }
   // theap->cookie = _mi_theap_random_next(theap) | 1;
   _mi_theap_guarded_init(theap); // needs theap->random
+  #if MI_PROFILE
+  // A profiler that is running samples this theap from its first allocation on, and not only after
+  // `page.c:mi_malloc_generic_admin` has looked at it (every 1000 generic allocations): a thread that 
+  // makes a few large allocations and ends would never be sampled.
+  if (!theap->is_detached) {
+    mi_profiler_t* const prof = mi_atomic_load_ptr_acquire(mi_profiler_t,&heap->profiler);
+    if (prof!=NULL && mi_profiler_is_enabled(prof) && prof->on_alloc!=NULL) {
+      _mi_theap_set_profile_sample_rate(theap, (prof->initial_sample_rate==0 ? 1 : prof->initial_sample_rate));
+      // and start with a full period: the countdowns are 0 here, and with that the first allocation would be
+      // a sample that reports a whole period of bytes that nobody requested.
+      theap->profile_sample_countdown = theap->profile_sample_rate;
+      theap->sample_countdown = theap->sample_rate;
+    }
+  }
+  #endif
   if (!theap->is_detached) {
     mi_subproc_stat_increase(_mi_theap_subproc(theap),theaps,1);  // on subproc to match theap_free_mem
   }