Reland: add dump provider for malloc heap profiler (https://codereview.chromium.org/1675183006/)

Reason for reland:
the previous CL forgot about the fact that nacl has its own special
base target (base_nacl) which is a fork of the real base one.
That caused a revert in crrev.com/1822013002 because:
 - The original CL introduced a dependency from a translation unit in
   base (malloc_dump_provider.cc) to allocator_features.h, which is a
   generated header.
 - The original CL was relying on the fact that there is an existing
   dependency from base to allocator_features declared in base.gyp.
 - The original CL forgot about the need of the same dependency from
   base_nacl.
This cl is fixing the latter point.

Original breakage:
https://build.chromium.org/p/chromium.chrome/builders/Google%20Chrome%20Mac/builds/8808

Original issue's description:
> 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
>
> Committed: https://crrev.com/14c8b52dac1285bbf23514d05e6109aa7edc427f
> Cr-Commit-Position: refs/heads/master@{#382505}

BUG=550886, 596873
TBR=ssid@chromium.org,petrcermak@chromium.org,dskiba@google.com,haraken@chromium.org

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

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

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);
     }
   }