tracing: add dump provider for malloc heap profiler

Now that the allocator shim is landing (the base API landed in
crre.com/1675143004, Linux support in crrev.com/1781573002) we
can use that to implement the malloc heap profiler on top of that.
This CL leverages the new shim API to collect information about
malloc/free and injecting that into the heap profiler, leveraging
the same infrastructure we are already using for PartitionAlloc
and BlinkGC (oilpan).

Next steps
----------
Currently this heap profiler reports also the memory used by the
tracing subsystem itself, which is not desirable. Will come up
with some scoped suppression semantic to bypass the reporting.

BUG=550886
TBR=haraken

Review URL: https://codereview.chromium.org/1675183006

Cr-Commit-Position: refs/heads/master@{#382505}

6 files changed, 232 insertions, 33 deletions

diff --git a/base/trace_event/heap_profiler_allocation_context_tracker.cc b/base/trace_event/heap_profiler_allocation_context_tracker.cc
index 791ab7a..2877d2a 100644
--- a/base/trace_event/heap_profiler_allocation_context_tracker.cc
+++ b/base/trace_event/heap_profiler_allocation_context_tracker.cc
@@ -18,6 +18,10 @@ subtle::Atomic32 AllocationContextTracker::capture_enabled_ = 0;
 
 namespace {
 
+const size_t kMaxStackDepth = 128u;
+AllocationContextTracker* const kInitializingSentinel =
+    reinterpret_cast<AllocationContextTracker*>(-1);
+
 ThreadLocalStorage::StaticSlot g_tls_alloc_ctx_tracker = TLS_INITIALIZER;
 
 // This function is added to the TLS slot to clean up the instance when the
@@ -28,15 +32,16 @@ void DestructAllocationContextTracker(void* alloc_ctx_tracker) {
 
 }  // namespace
 
-AllocationContextTracker::AllocationContextTracker() {}
-AllocationContextTracker::~AllocationContextTracker() {}
-
 // static
-AllocationContextTracker* AllocationContextTracker::GetThreadLocalTracker() {
-  auto tracker =
+AllocationContextTracker*
+AllocationContextTracker::GetInstanceForCurrentThread() {
+  AllocationContextTracker* tracker =
       static_cast<AllocationContextTracker*>(g_tls_alloc_ctx_tracker.Get());
+  if (tracker == kInitializingSentinel)
+    return nullptr;  // Re-entrancy case.
 
   if (!tracker) {
+    g_tls_alloc_ctx_tracker.Set(kInitializingSentinel);
     tracker = new AllocationContextTracker();
     g_tls_alloc_ctx_tracker.Set(tracker);
   }
@@ -44,6 +49,11 @@ AllocationContextTracker* AllocationContextTracker::GetThreadLocalTracker() {
   return tracker;
 }
 
+AllocationContextTracker::AllocationContextTracker() {
+  pseudo_stack_.reserve(kMaxStackDepth);
+}
+AllocationContextTracker::~AllocationContextTracker() {}
+
 // static
 void AllocationContextTracker::SetCaptureEnabled(bool enabled) {
   // When enabling capturing, also initialize the TLS slot. This does not create
@@ -56,45 +66,41 @@ void AllocationContextTracker::SetCaptureEnabled(bool enabled) {
   subtle::Release_Store(&capture_enabled_, enabled);
 }
 
-// static
 void AllocationContextTracker::PushPseudoStackFrame(StackFrame frame) {
-  auto tracker = AllocationContextTracker::GetThreadLocalTracker();
-
   // Impose a limit on the height to verify that every push is popped, because
   // in practice the pseudo stack never grows higher than ~20 frames.
-  DCHECK_LT(tracker->pseudo_stack_.size(), 128u);
-  tracker->pseudo_stack_.push_back(frame);
+  if (pseudo_stack_.size() < kMaxStackDepth)
+    pseudo_stack_.push_back(frame);
+  else
+    NOTREACHED();
 }
 
 // static
 void AllocationContextTracker::PopPseudoStackFrame(StackFrame frame) {
-  auto tracker = AllocationContextTracker::GetThreadLocalTracker();
-
   // Guard for stack underflow. If tracing was started with a TRACE_EVENT in
   // scope, the frame was never pushed, so it is possible that pop is called
   // on an empty stack.
-  if (tracker->pseudo_stack_.empty())
+  if (pseudo_stack_.empty())
     return;
 
   // Assert that pushes and pops are nested correctly. This DCHECK can be
   // hit if some TRACE_EVENT macro is unbalanced (a TRACE_EVENT_END* call
   // without a corresponding TRACE_EVENT_BEGIN).
-  DCHECK_EQ(frame, tracker->pseudo_stack_.back())
+  DCHECK_EQ(frame, pseudo_stack_.back())
       << "Encountered an unmatched TRACE_EVENT_END";
 
-  tracker->pseudo_stack_.pop_back();
+  pseudo_stack_.pop_back();
 }
 
 // static
 AllocationContext AllocationContextTracker::GetContextSnapshot() {
-  AllocationContextTracker* tracker = GetThreadLocalTracker();
   AllocationContext ctx;
 
   // Fill the backtrace.
   {
-    auto src = tracker->pseudo_stack_.begin();
+    auto src = pseudo_stack_.begin();
     auto dst = std::begin(ctx.backtrace.frames);
-    auto src_end = tracker->pseudo_stack_.end();
+    auto src_end = pseudo_stack_.end();
     auto dst_end = std::end(ctx.backtrace.frames);
 
     // Copy as much of the bottom of the pseudo stack into the backtrace as
diff --git a/base/trace_event/heap_profiler_allocation_context_tracker.h b/base/trace_event/heap_profiler_allocation_context_tracker.h
index 9c9a313..20a6e30 100644
--- a/base/trace_event/heap_profiler_allocation_context_tracker.h
+++ b/base/trace_event/heap_profiler_allocation_context_tracker.h
@@ -43,22 +43,25 @@ class BASE_EXPORT AllocationContextTracker {
     return subtle::Acquire_Load(&capture_enabled_) != 0;
   }
 
+  // Returns the thread-local instance, creating one if necessary. Returns
+  // always a valid instance, unless it is called re-entrantly, in which case
+  // returns nullptr in the nested calls.
+  static AllocationContextTracker* GetInstanceForCurrentThread();
+
   // Pushes a frame onto the thread-local pseudo stack.
-  static void PushPseudoStackFrame(StackFrame frame);
+  void PushPseudoStackFrame(StackFrame frame);
 
   // Pops a frame from the thread-local pseudo stack.
-  static void PopPseudoStackFrame(StackFrame frame);
+  void PopPseudoStackFrame(StackFrame frame);
 
   // Returns a snapshot of the current thread-local context.
-  static AllocationContext GetContextSnapshot();
+  AllocationContext GetContextSnapshot();
 
   ~AllocationContextTracker();
 
  private:
   AllocationContextTracker();
 
-  static AllocationContextTracker* GetThreadLocalTracker();
-
   static subtle::Atomic32 capture_enabled_;
 
   // The pseudo stack where frames are |TRACE_EVENT| names.
diff --git a/base/trace_event/heap_profiler_allocation_context_tracker_unittest.cc b/base/trace_event/heap_profiler_allocation_context_tracker_unittest.cc
index 58255ad..2498fcd 100644
--- a/base/trace_event/heap_profiler_allocation_context_tracker_unittest.cc
+++ b/base/trace_event/heap_profiler_allocation_context_tracker_unittest.cc
@@ -27,7 +27,9 @@ const char kGingerbread[] = "Gingerbread";
 // in |AllocationContextTracker::GetContextSnapshot|.
 template <size_t N>
 void AssertBacktraceEquals(const StackFrame(&expected_backtrace)[N]) {
-  AllocationContext ctx = AllocationContextTracker::GetContextSnapshot();
+  AllocationContext ctx =
+      AllocationContextTracker::GetInstanceForCurrentThread()
+          ->GetContextSnapshot();
 
   auto actual = std::begin(ctx.backtrace.frames);
   auto actual_bottom = std::end(ctx.backtrace.frames);
@@ -46,7 +48,9 @@ void AssertBacktraceEquals(const StackFrame(&expected_backtrace)[N]) {
 }
 
 void AssertBacktraceEmpty() {
-  AllocationContext ctx = AllocationContextTracker::GetContextSnapshot();
+  AllocationContext ctx =
+      AllocationContextTracker::GetInstanceForCurrentThread()
+          ->GetContextSnapshot();
 
   for (StackFrame frame : ctx.backtrace.frames)
     ASSERT_EQ(nullptr, frame);
@@ -207,7 +211,9 @@ TEST_F(AllocationContextTrackerTest, BacktraceTakesTop) {
 
   {
     TRACE_EVENT0("Testing", kGingerbread);
-    AllocationContext ctx = AllocationContextTracker::GetContextSnapshot();
+    AllocationContext ctx =
+        AllocationContextTracker::GetInstanceForCurrentThread()
+            ->GetContextSnapshot();
 
     // The pseudo stack relies on pointer equality, not deep string comparisons.
     ASSERT_EQ(kCupcake, ctx.backtrace.frames[0]);
@@ -215,7 +221,9 @@ TEST_F(AllocationContextTrackerTest, BacktraceTakesTop) {
   }
 
   {
-    AllocationContext ctx = AllocationContextTracker::GetContextSnapshot();
+    AllocationContext ctx =
+        AllocationContextTracker::GetInstanceForCurrentThread()
+            ->GetContextSnapshot();
     ASSERT_EQ(kCupcake, ctx.backtrace.frames[0]);
     ASSERT_EQ(kFroyo, ctx.backtrace.frames[11]);
   }
diff --git a/base/trace_event/malloc_dump_provider.cc b/base/trace_event/malloc_dump_provider.cc
index 6f9aa96..d8f82ed 100644
--- a/base/trace_event/malloc_dump_provider.cc
+++ b/base/trace_event/malloc_dump_provider.cc
@@ -7,7 +7,14 @@
 #include <stddef.h>
 
 #include "base/allocator/allocator_extension.h"
+#include "base/allocator/allocator_shim.h"
+#include "base/allocator/features.h"
+#include "base/trace_event/heap_profiler_allocation_context.h"
+#include "base/trace_event/heap_profiler_allocation_context_tracker.h"
+#include "base/trace_event/heap_profiler_allocation_register.h"
+#include "base/trace_event/heap_profiler_heap_dump_writer.h"
 #include "base/trace_event/process_memory_dump.h"
+#include "base/trace_event/trace_event_argument.h"
 #include "build/build_config.h"
 
 #if defined(OS_MACOSX)
@@ -19,6 +26,65 @@
 namespace base {
 namespace trace_event {
 
+#if BUILDFLAG(USE_EXPERIMENTAL_ALLOCATOR_SHIM)
+namespace {
+
+using allocator::AllocatorDispatch;
+
+void* HookAlloc(const AllocatorDispatch* self, size_t size) {
+  const AllocatorDispatch* const next = self->next;
+  void* ptr = next->alloc_function(next, size);
+  if (ptr)
+    MallocDumpProvider::GetInstance()->InsertAllocation(ptr, size);
+  return ptr;
+}
+
+void* HookZeroInitAlloc(const AllocatorDispatch* self, size_t n, size_t size) {
+  const AllocatorDispatch* const next = self->next;
+  void* ptr = next->alloc_zero_initialized_function(next, n, size);
+  if (ptr)
+    MallocDumpProvider::GetInstance()->InsertAllocation(ptr, n * size);
+  return ptr;
+}
+
+void* HookllocAligned(const AllocatorDispatch* self,
+                      size_t alignment,
+                      size_t size) {
+  const AllocatorDispatch* const next = self->next;
+  void* ptr = next->alloc_aligned_function(next, alignment, size);
+  if (ptr)
+    MallocDumpProvider::GetInstance()->InsertAllocation(ptr, size);
+  return ptr;
+}
+
+void* HookRealloc(const AllocatorDispatch* self, void* address, size_t size) {
+  const AllocatorDispatch* const next = self->next;
+  void* ptr = next->realloc_function(next, address, size);
+  MallocDumpProvider::GetInstance()->RemoveAllocation(address);
+  if (size > 0)  // realloc(size == 0) means free().
+    MallocDumpProvider::GetInstance()->InsertAllocation(ptr, size);
+  return ptr;
+}
+
+void HookFree(const AllocatorDispatch* self, void* address) {
+  if (address)
+    MallocDumpProvider::GetInstance()->RemoveAllocation(address);
+  const AllocatorDispatch* const next = self->next;
+  next->free_function(next, address);
+}
+
+AllocatorDispatch g_allocator_hooks = {
+    &HookAlloc,         /* alloc_function */
+    &HookZeroInitAlloc, /* alloc_zero_initialized_function */
+    &HookllocAligned,   /* alloc_aligned_function */
+    &HookRealloc,       /* realloc_function */
+    &HookFree,          /* free_function */
+    nullptr,            /* next */
+};
+
+}  // namespace
+#endif  // BUILDFLAG(USE_EXPERIMENTAL_ALLOCATOR_SHIM)
+
 // static
 const char MallocDumpProvider::kAllocatedObjects[] = "malloc/allocated_objects";
 
@@ -28,7 +94,8 @@ MallocDumpProvider* MallocDumpProvider::GetInstance() {
                    LeakySingletonTraits<MallocDumpProvider>>::get();
 }
 
-MallocDumpProvider::MallocDumpProvider() {}
+MallocDumpProvider::MallocDumpProvider()
+    : heap_profiler_enabled_(false), tid_dumping_heap_(kInvalidThreadId) {}
 
 MallocDumpProvider::~MallocDumpProvider() {}
 
@@ -96,8 +163,99 @@ bool MallocDumpProvider::OnMemoryDump(const MemoryDumpArgs& args,
                           resident_size - allocated_objects_size);
   }
 
+  // Heap profiler dumps.
+  if (!heap_profiler_enabled_)
+    return true;
+
+  // The dumps of the heap profiler should be created only when heap profiling
+  // was enabled (--enable-heap-profiling) AND a DETAILED dump is requested.
+  // However, when enabled, the overhead of the heap profiler should be always
+  // reported to avoid oscillations of the malloc total in LIGHT dumps.
+
+  tid_dumping_heap_ = PlatformThread::CurrentId();
+  // At this point the Insert/RemoveAllocation hooks will ignore this thread.
+  // Enclosing all the temporariy data structures in a scope, so that the heap
+  // profiler does not see unabalanced malloc/free calls from these containers.
+  {
+    TraceEventMemoryOverhead overhead;
+    hash_map<AllocationContext, size_t> bytes_by_context;
+    {
+      AutoLock lock(allocation_register_lock_);
+      if (allocation_register_) {
+        if (args.level_of_detail == MemoryDumpLevelOfDetail::DETAILED) {
+          for (const auto& alloc_size : *allocation_register_)
+            bytes_by_context[alloc_size.context] += alloc_size.size;
+        }
+        allocation_register_->EstimateTraceMemoryOverhead(&overhead);
+      }
+    }  // lock(allocation_register_lock_)
+
+    if (!bytes_by_context.empty()) {
+      scoped_ptr<TracedValue> heap_dump = ExportHeapDump(
+          bytes_by_context, pmd->session_state()->stack_frame_deduplicator(),
+          pmd->session_state()->type_name_deduplicator());
+      pmd->AddHeapDump("malloc", std::move(heap_dump));
+    }
+    overhead.DumpInto("tracing/heap_profiler_malloc", pmd);
+  }
+  tid_dumping_heap_ = kInvalidThreadId;
+
   return true;
 }
 
+void MallocDumpProvider::OnHeapProfilingEnabled(bool enabled) {
+#if BUILDFLAG(USE_EXPERIMENTAL_ALLOCATOR_SHIM)
+  if (enabled) {
+    {
+      AutoLock lock(allocation_register_lock_);
+      allocation_register_.reset(new AllocationRegister());
+    }
+    allocator::InsertAllocatorDispatch(&g_allocator_hooks);
+  } else {
+    AutoLock lock(allocation_register_lock_);
+    allocation_register_.reset();
+    // Insert/RemoveAllocation below will no-op if the register is torn down.
+    // Once disabled, heap profiling will not re-enabled anymore for the
+    // lifetime of the process.
+  }
+#endif
+  heap_profiler_enabled_ = enabled;
+}
+
+void MallocDumpProvider::InsertAllocation(void* address, size_t size) {
+  // CurrentId() can be a slow operation (crbug.com/497226). This apparently
+  // redundant condition short circuits the CurrentID() calls when unnecessary.
+  if (tid_dumping_heap_ != kInvalidThreadId &&
+      tid_dumping_heap_ == PlatformThread::CurrentId())
+    return;
+
+  // AllocationContextTracker will return nullptr when called re-reentrantly.
+  // This is the case of GetInstanceForCurrentThread() being called for the
+  // first time, which causes a new() inside the tracker which re-enters the
+  // heap profiler, in which case we just want to early out.
+  auto tracker = AllocationContextTracker::GetInstanceForCurrentThread();
+  if (!tracker)
+    return;
+  AllocationContext context = tracker->GetContextSnapshot();
+
+  AutoLock lock(allocation_register_lock_);
+  if (!allocation_register_)
+    return;
+
+  allocation_register_->Insert(address, size, context);
+}
+
+void MallocDumpProvider::RemoveAllocation(void* address) {
+  // No re-entrancy is expected here as none of the calls below should
+  // cause a free()-s (|allocation_register_| does its own heap management).
+  if (tid_dumping_heap_ != kInvalidThreadId &&
+      tid_dumping_heap_ == PlatformThread::CurrentId())
+    return;
+  AutoLock lock(allocation_register_lock_);
+  if (!allocation_register_)
+    return;
+  allocation_register_->Remove(address);
+}
+
 }  // namespace trace_event
 }  // namespace base
diff --git a/base/trace_event/malloc_dump_provider.h b/base/trace_event/malloc_dump_provider.h
index 63fc1b0..d9d7021 100644
--- a/base/trace_event/malloc_dump_provider.h
+++ b/base/trace_event/malloc_dump_provider.h
@@ -8,7 +8,10 @@
 #include <istream>
 
 #include "base/macros.h"
+#include "base/memory/scoped_ptr.h"
 #include "base/memory/singleton.h"
+#include "base/synchronization/lock.h"
+#include "base/threading/platform_thread.h"
 #include "base/trace_event/memory_dump_provider.h"
 #include "build/build_config.h"
 
@@ -20,6 +23,8 @@
 namespace base {
 namespace trace_event {
 
+class AllocationRegister;
+
 // Dump provider which collects process-wide memory stats.
 class BASE_EXPORT MallocDumpProvider : public MemoryDumpProvider {
  public:
@@ -33,12 +38,28 @@ class BASE_EXPORT MallocDumpProvider : public MemoryDumpProvider {
   bool OnMemoryDump(const MemoryDumpArgs& args,
                     ProcessMemoryDump* pmd) override;
 
+  void OnHeapProfilingEnabled(bool enabled) override;
+
+  // For heap profiling.
+  void InsertAllocation(void* address, size_t size);
+  void RemoveAllocation(void* address);
+
  private:
   friend struct DefaultSingletonTraits<MallocDumpProvider>;
 
   MallocDumpProvider();
   ~MallocDumpProvider() override;
 
+  // For heap profiling.
+  bool heap_profiler_enabled_;
+  scoped_ptr<AllocationRegister> allocation_register_;
+  Lock allocation_register_lock_;
+
+  // When in OnMemoryDump(), this contains the current thread ID.
+  // This is to prevent re-entrancy in the heap profiler when the heap dump
+  // generation is malloc/new-ing for its own bookeeping data structures.
+  PlatformThreadId tid_dumping_heap_;
+
   DISALLOW_COPY_AND_ASSIGN(MallocDumpProvider);
 };
 
diff --git a/base/trace_event/trace_log.cc b/base/trace_event/trace_log.cc
index dcedb44..3b1d2c9 100644
--- a/base/trace_event/trace_log.cc
+++ b/base/trace_event/trace_log.cc
@@ -1318,12 +1318,14 @@ TraceEventHandle TraceLog::AddTraceEventWithThreadIdAndTimestamp(
   if (!(flags & TRACE_EVENT_FLAG_COPY)) {
     if (AllocationContextTracker::capture_enabled()) {
       if (phase == TRACE_EVENT_PHASE_BEGIN ||
-          phase == TRACE_EVENT_PHASE_COMPLETE)
-        AllocationContextTracker::PushPseudoStackFrame(name);
-      else if (phase == TRACE_EVENT_PHASE_END)
+          phase == TRACE_EVENT_PHASE_COMPLETE) {
+        AllocationContextTracker::GetInstanceForCurrentThread()
+            ->PushPseudoStackFrame(name);
+      } else if (phase == TRACE_EVENT_PHASE_END)
         // The pop for |TRACE_EVENT_PHASE_COMPLETE| events
         // is in |TraceLog::UpdateTraceEventDuration|.
-        AllocationContextTracker::PopPseudoStackFrame(name);
+        AllocationContextTracker::GetInstanceForCurrentThread()
+            ->PopPseudoStackFrame(name);
     }
   }
 
@@ -1447,7 +1449,8 @@ void TraceLog::UpdateTraceEventDuration(
 
     if (base::trace_event::AllocationContextTracker::capture_enabled()) {
       // The corresponding push is in |AddTraceEventWithThreadIdAndTimestamp|.
-      base::trace_event::AllocationContextTracker::PopPseudoStackFrame(name);
+      base::trace_event::AllocationContextTracker::GetInstanceForCurrentThread()
+          ->PopPseudoStackFrame(name);
     }
   }