Check populated before calling unpopulate in ReleaseSlabMetadataForDrainedCpus.

docs/sampling.md

@@ -43,10 +43,13 @@ We also tell the span that we're sampling it. We can do this because we do
 sampling at tcmalloc page sizes, so each sample corresponds to a particular page
 in the pagemap.
 
-For small allocations, the returned allocation uses an entire TCMalloc page (not
-shared with any other allocations) on a dedicated span. The returned allocation
-is placed on the sampled page heap, allowing us to use the pointer's tag bits to
-identify that the object was sampled and needs special handling on deallocation.
+For small allocations, we make up to two allocations: the returned allocation
+(which uses an entire TCMalloc page, not shared with any other allocations) and
+a proxy allocation in a non-sampled span (the proxy object was formerly used for
+computing fragmentation profiles) for sizes with >1 objects-per-span (mostly
+sizes <8KB in the default configuration). The returned allocation is placed on
+the sampled page heap, allowing us to use the pointer's tag bits to identify
+that the object was sampled and needs special handling on deallocation.
 
 For the sampled page heap, the virtual addresses associated with the allocation
 are
@@ -55,11 +58,11 @@ This, combined with the whole-page behavior above, means that *every allocation
 gets its own native (OS) page(s)* shared with no other allocations.
 
 For large (`>kMaxSize`) allocations, the returned allocation will be on entire
-TCMalloc pages. These objects are requested directly from the non-sampled page
-heaps. These objects will be packed by [Temeraire](temeraire.md) densely onto
-hugepages. While objects that are exact multiples of 2MB are given their own
-hugepages (the `HugeCache`), the access patterns of other objects may affect
-statistics for the profiled ones.
+TCMalloc pages and there is no proxy object. These objects are requested
+directly from the non-sampled page heaps. These objects will be packed by
+[Temeraire](temeraire.md) densely onto hugepages. While objects that are exact
+multiples of 2MB are given their own hugepages (the `HugeCache`), the access
+patterns of other objects may affect statistics for the profiled ones.
 
 | Statistic | Small                         | Large                        |
 | :-------- | :---------------------------- | :--------------------------- |
@@ -75,10 +78,11 @@ we can quickly test whether a particular allocation might be a sample. When we
 are done with the sampled span we release it using
 [tcmalloc::Span::Unsample()](https://github.com/google/tcmalloc/blob/master/tcmalloc/span.cc).
 
-## How Do We Handle Heap Profiling
+## How Do We Handle Heap and Fragmentation Profiling
 
-To handle heap profiling we just need to traverse the list of sampled objects
-and compute the amount of heap they consume.
+To handle heap and fragmentation profiling we just need to traverse the list of
+sampled objects and compute either their degree of fragmentation (with the proxy
+object), or the amount of heap they consume.
 
 Each allocation gets additional metadata associated with it when it is exposed
 in the heap profile. In the preparation for writing the heap profile,

tcmalloc/huge_page_tracker.h

@@ -232,27 +232,27 @@ class PageTracker : public TList<PageTracker>::Elem {
   void SetHasDenseSpans() { has_dense_spans_ = true; }
 
   struct HugePageResidencyState {
-    // Records the unbacked bitmap for this hugepage. In terms of TCMalloc
-    // pages. scaled via `ReductionOp::kAll`.
-    PageBitmap unbacked;
-    // Records the swapped bitmap for this hugepage. In terms of TCMalloc
-    // pages. scaled via `ReductionOp::kAny`.
-    PageBitmap swapped;
-    // Records the stale bitmap for this hugepage. In terms of TCMalloc
-    // pages. scaled via `ReductionOp::kAny`.
-    PageBitmap stale;
+    // Records whether the page is hugepage backed.
+    bool maybe_hugepage_backed = false;
     // Records the time (in ticks) when the residency state was last updated.
     // This is used to determine when the tracker may be revisited for
     // collapse.
     double record_time;
-    // Records whether the page is hugepage backed.
-    bool maybe_hugepage_backed = false;
     // Records whether metrics are valid. It is set the first time the
     // residency state is queried.
     bool entry_valid = false;
     // This records the trackers that are currently being collapsed. This is
     // used to avoid subreleasing the pages that are being collapsed.
     bool being_collapsed = false;
+    // Records the unbacked bitmap for this hugepage. In terms of TCMalloc
+    // pages. scaled via `ReductionOp::kAll`.
+    PageBitmap unbacked;
+    // Records the swapped bitmap for this hugepage. In terms of TCMalloc
+    // pages. scaled via `ReductionOp::kAny`.
+    PageBitmap swapped;
+    // Records the stale bitmap for this hugepage. In terms of TCMalloc
+    // pages. scaled via `ReductionOp::kAny`.
+    PageBitmap stale;
     // Records whether collapse was skipped due to threshold constraints.
     bool collapse_skipped = false;
     // Records whether collapse was skipped due to backoff.
@@ -350,13 +350,6 @@ class PageTracker : public TList<PageTracker>::Elem {
   // reset it once we measure those pages in abandoned_count_.
   bool abandoned_;
   bool unbroken_;
-  bool has_dense_spans_ = false;
-  // This field is used to avoid freeing this tracker prematurely. When this
-  // is set, any maintenance operation (e.g. collapse) that drops
-  // pageheap_lock might manipulate the tracker state without holding the
-  // lock. When all the pages on the tracked hugepage are freed, this field
-  // is checked to ensure that the tracker is not freed right away.
-  uint8_t dont_free_tracker_mask_ = 0;
   double alloctime_;
   double last_page_allocation_time_ = 0;
 
@@ -385,8 +378,17 @@ class PageTracker : public TList<PageTracker>::Elem {
                     std::numeric_limits<uint16_t>::max(),
                 "nallocs must be able to support kPagesPerHugePage!");
 
+  bool has_dense_spans_ = false;
+
   HugePageResidencyState hugepage_residency_state_;
 
+  // This field is used to avoid freeing this tracker prematurely. When this
+  // is set, any maintenance operation (e.g. collapse) that drops
+  // pageheap_lock might manipulate the tracker state without holding the
+  // lock. When all the pages on the tracked hugepage are freed, this field
+  // is checked to ensure that the tracker is not freed right away.
+  uint8_t dont_free_tracker_mask_ = 0;
+
   [[nodiscard]] bool ReleasePages(Range r, MemoryModifyFunction& unback) {
     bool success = unback(r).success;
     if (ABSL_PREDICT_TRUE(success)) {
@@ -437,21 +439,31 @@ inline PageTracker::HardwarePageResidencyInfo PageTracker::CountInfoInHugePage(
   }
   TC_ASSERT_LE(kHardwarePagesInHugePage, kMaxResidencyBits);
 
-  const PageBitmap& used = free_.bits();
-  const PageBitmap free = ~used;
+  const PageBitmap& free = free_.bits();
 
   TC_ASSERT_EQ(kHardwarePagesInHugePage % kPagesPerHugePage.raw_num(), 0);
   const int shift = kHardwarePagesInHugePage / kPagesPerHugePage.raw_num();
   const int shift_bits = absl::bit_width<uint8_t>(shift - 1);
   TC_ASSERT_LT((kHardwarePagesInHugePage - 1) >> shift_bits,
                kPagesPerHugePage.raw_num());
 
-  return {.n_free_swapped = (free & swapped).CountBits() * shift,
-          .n_used_swapped = (used & swapped).CountBits() * shift,
-          .n_free_unbacked = (free & unbacked).CountBits() * shift,
-          .n_used_unbacked = (used & unbacked).CountBits() * shift,
-          .n_free_stale = (free & stale).CountBits() * shift,
-          .n_used_stale = (used & stale).CountBits() * shift};
+  size_t n_unbacked[2] = {0, 0};
+  size_t n_swapped[2] = {0, 0};
+  size_t n_stale[2] = {0, 0};
+
+  n_unbacked[0] = (free & unbacked).CountBits() * shift;
+  n_unbacked[1] = (~free & unbacked).CountBits() * shift;
+  n_swapped[0] = (free & swapped).CountBits() * shift;
+  n_swapped[1] = (~free & swapped).CountBits() * shift;
+  n_stale[0] = (free & stale).CountBits() * shift;
+  n_stale[1] = (~free & stale).CountBits() * shift;
+
+  return {.n_free_swapped = n_swapped[1],
+          .n_used_swapped = n_swapped[0],
+          .n_free_unbacked = n_unbacked[1],
+          .n_used_unbacked = n_unbacked[0],
+          .n_free_stale = n_stale[1],
+          .n_used_stale = n_stale[0]};
 }
 
 inline void PageTracker::Put(Range r, SpanAllocInfo span_alloc_info) {

tcmalloc/internal/percpu_tcmalloc.h

@@ -1543,7 +1543,9 @@ void TcmallocSlab<NumClasses>::ReleaseSlabMetadataForDrainedCpus(
                         reinterpret_cast<uintptr_t>(slabs)) /
                        slab_size_bytes;
     for (unsigned i = 0; i < bytes_to_free / slab_size_bytes; ++i) {
-      unpopulate(first_cpu + i);
+      if (populated(first_cpu + i)) {
+        unpopulate(first_cpu + i);
+      }
     }
   }
 

tcmalloc/span.h

@@ -254,10 +254,6 @@ class ABSL_CACHELINE_ALIGNED Span final : public SpanList::Elem {
   [[nodiscard]] ObjIdx BitmapPtrToIdx(void* ptr, size_t size,
                                       uint32_t reciprocal) const;
   [[nodiscard]] void* BitmapIdxToPtr(ObjIdx idx, size_t size) const;
-#ifndef TCMALLOC_INTERNAL_LEGACY_LOCKING
-  [[nodiscard]] void* BitmapIdxToPtr(ObjIdx idx, size_t size,
-                                     uintptr_t start) const;
-#endif
 
 #ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING
   static constexpr size_t kNonemptyIndexBits = 5;
@@ -610,15 +606,9 @@ inline Span::ObjIdx Span::OffsetToIdx(uintptr_t offset, uint32_t reciprocal) {
 }
 
 #ifndef TCMALLOC_INTERNAL_LEGACY_LOCKING
-inline void* Span::BitmapIdxToPtr(ObjIdx idx, size_t size,
-                                  uintptr_t start) const {
-  TC_ASSERT_EQ(start, first_page().start_uintptr());
-  uintptr_t off = start + idx * size;
-  return reinterpret_cast<void*>(off);
-}
-
 inline void* Span::BitmapIdxToPtr(ObjIdx idx, size_t size) const {
-  return BitmapIdxToPtr(idx, size, first_page().start_uintptr());
+  uintptr_t off = first_page().start_uintptr() + idx * size;
+  return reinterpret_cast<void*>(off);
 }
 #endif
 
@@ -780,10 +770,9 @@ inline size_t Span::BitmapPopBatch(absl::Span<void*> batch,
   return count;
 #else
   void** ptrs = batch.data();
-  const uintptr_t span_start = first_page().start_uintptr();
   size_t popped = bitmap_.PopBatch(
       [&](size_t offset) {
-        *ptrs++ = BitmapIdxToPtr(static_cast<ObjIdx>(offset), size, span_start);
+        *ptrs++ = BitmapIdxToPtr(static_cast<ObjIdx>(offset), size);
       },
       batch.size());
   allocated_.store(allocated_.load(std::memory_order_relaxed) + popped,