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

[Jarred-Sumner]

Five changes to the sampled-profiling hooks (mi_profiler_t) that Bun's heap profiler needs. None of them adds an instruction to the mi_malloc/mi_free fast paths. Each comes with a test in test/test-profile.c that fails at ab135013.

1. A small allocation is sampled when its own size class used up the countdown (page.c)

With coarse sampling (MI_SAMPLE==1, what MI_PROFILE=1 gives) the fast path does not touch the sample countdown. The blocks it took from a page are counted when the page's free list is refilled in the generic path. _mi_malloc_generic looked at the countdown before mi_page_queue_find_free did that refill, so when the refill used up the countdown the sample went to the next allocation to come through the generic path. Every allocation above MI_SMALL_SIZE_MAX comes through it, so those took nearly all the samples that small allocations had earned.

test_profiler_small_attribution allocates 40 blocks of 64 bytes for each block of 1100 bytes (70% of the bytes are in the 64-byte blocks) and keeps them all. Share of the sampled bytes that went to the 64-byte blocks:

	share (0.699 is exact)
ab135013, MI_PROFILE=1	0.081
this branch, MI_PROFILE=1	0.585
MI_PROFILE=2 (exact countdown in the fast path)	0.681

The change: after mi_page_queue_find_free, if sampling is on and the countdown is used up, the allocation in progress is the sample (it has the size class of the blocks that were just counted). mi_page_extend_free now counts the blocks of the previous extension too; before, a page that only grew (nothing freed) was first counted when it became full.

What remains of the bias comes from the counting being late by up to one refill (4 KiB); it goes either way depending on the mix and is bounded by that.

Cost when no profiler runs: a load and a branch on theap->sample_rate per call of _mi_malloc_generic that finds a page. On the malloc/free benchmark of #37 (25 M mallocs, 25 M frees; clang 21 -O3 -march=nehalem, as C++): 2.26451 G -> 2.26696 G user instructions (+0.11%). mi_malloc and mi_free are byte-identical. (I tried folding it into the existing check in front of mi_page_queue_find_free, which is one check instead of two: clang then keeps page live across it and that costs more, +0.29%.)

2. bytes_since_last_sample stays exact when on_alloc changes the rate (sample-profile.c)

_mi_theap_malloc_sampled takes the period that just ended to have been sample_rate bytes long. When on_alloc returned a higher rate, _mi_theap_set_profile_sample_rate left the countdown at the old rate ("todo: adjust difference?"), so the next period was the old length but was reported with the new one: each sample reported the difference on top. With a rate that varies around a mean (which is how a profiler avoids sampling a periodic pattern always at the same allocation) the samples added up to 18% more than was allocated in the test, and up to 50% with an exponential distribution. Now the new period starts in full at a rate change. The same test checks that the samples add up to what was requested.

3. A thread that starts while a profiler runs is sampled from its first allocation (theap.c)

A theap only noticed an enabled profiler in mi_malloc_generic_admin, every 1000 generic allocations. A thread that makes a hundred large allocations and ends was never sampled; a new heap (an arena) on a thread lost its first ~4 MiB. _mi_theap_init now sets the profile rate when the heap's profiler is already enabled, and starts with a full period (the countdowns are 0 at that point: without that the first allocation of each new theap is a sample that reports initial_sample_rate bytes nobody requested). test_profiler_new_thread: 8 allocations of 1 MiB on a new thread, 8 samples or more (was 1). test_profiler_new_thread_no_phantom_period: 4 threads that allocate 32 bytes each report 0 bytes (2 MiB without the full period, with initial_sample_rate = 512 KiB).

4. mi_heap_destroy reports the sampled blocks it frees (sample-profile.c, arena.c)

mi_heap_destroy frees pages without a mi_free per block, so on_free was never called for the sampled blocks in them and a profiler counted them as in use forever. There is _mi_page_unguard_all at that spot for guarded blocks; _mi_page_profile_free_all next to it walks the used blocks of a page that has the interior-pointer flag (only pages that ever held a sampled block do) and calls on_free for the profiled ones.

mi_free tells a profiled block by its tag and by being handed an interior pointer. The walk has no pointer to go by, and the first word of a regular block in use is program data that can hold the value of the tag. So a profiled block now carries a check word behind the tag (its address xor a random key), which _mi_page_profile_on_free clears (the free list only overwrites the tag, and the block can come back as a regular block). test_profiler_heap_destroy fills every block with {1, ...} and checks that exactly the sampled blocks are reported, once (those freed with mi_free before the destroy are not reported again). test_profiler_heap_destroy_recycled frees 400 sampled blocks, takes the same blocks again as regular ones with only the first word written (1), and destroys the heap: nothing is reported for them (12 of them were, in a release build, before the word was cleared; a debug build fills freed memory and hides it).

Layout of a profiled block: [tag][check][mi_profiler_sample_data_t ...][user block]; one word more than before.

Testing

Linux x64, the configurations of .github/workflows/test.yaml with their flags and environment: Release, Debug MI_DEBUG=FULL, Secure, Debug and Release as C++, ASAN, UBSAN, TSAN, Debug and Release MI_GUARDED=FULL (with the job's MIMALLOC_GUARDED_SAMPLE_RATE), Debug and Release MI_FREE_USE_PAGEMAP=ON (what Bun builds with), Debug and Release MI_SECURE=FULL, TLS pthreads, SIMD, and Release MI_PROFILE=OFF. ctest is 100% in each except test-emulated-tls in the secure and guarded ones (see below; the same without this branch). src/static.c compiles as C++20 with Bun's flags. At ab135013 the new tests fail (shares and counts above).

5. on_free gets the pointer that on_alloc got, also for an aligned allocation (sample-profile.c)

mi_malloc_aligned aligns inside the block that _mi_theap_malloc_profiled returned, so mi_free of a sampled aligned block called on_free with the aligned pointer and not the one on_alloc was given. test-profile.c's own on_free asserts they are the same; it only never met a sampled aligned block. It does with guarded sampling on in the same process (the CI's Debug, Guarded job sets MIMALLOC_GUARDED_SAMPLE_RATE for the whole job): samples then land on mimalloc's own mi_heap_zalloc_aligned of the per-heap arena page table, which mi_heap_destroy frees. _mi_page_profile_on_free now computes the pointer from the block's layout. It also checks the block's check word before it calls on_free (it did only in the mi_heap_destroy walk), so a block that merely starts with the tag's value and is freed through an interior pointer is not reported.

test_profiler_guarded_mixed: 6000 blocks of five sizes in a new heap, a third of them with an alignment (64, 1 KiB, 16 KiB), with guarded sampling on in guarded builds (it sets the rate itself; 1333 of the 6000 are guarded, 2801 sampled); a quarter is freed on this thread, a quarter on another thread, a quarter moved by mi_heap_realloc, the rest destroyed with the heap. Every sampled block is reported once with the pointer on_alloc got. Without this change it fails (the assert in debug builds, a wrong pointer in release builds).

Tests and guarded sampling

CMakeLists.txt deliberately gave test-profile no guarded sample rate, but the CI's guarded jobs set one in the environment for every test. With guarded samples in between, what the profiler's samples add up to is off at ab135013 already (34 of 70 MiB, share 0.88 for 0.70), so test-profile now gets MIMALLOC_GUARDED_SAMPLE_RATE=0 explicitly, like test-purge-holes; the case above turns guarded sampling on for itself.

Seen on the way, not changed

• Profile and guarded sampling in one theap: _mi_theap_malloc_sampled subtracts the bytes since the last profile sample from both derived countdowns at every combined sample, so a profile sample comes early and the reported bytes do not add up (the numbers above). Bun does not build with MI_GUARDED.
• test-emulated-tls fails here in the MI_SECURE and MI_GUARDED configurations with and without this branch ("unexpected block layout": it expects two consecutive small allocations to be adjacent); it passes in CI.