Record Chrome trace events in tcmalloc heap profiles

third_party/tcmalloc/chromium/src/heap-profiler.cc

Index: third_party/tcmalloc/chromium/src/heap-profiler.cc
diff --git a/third_party/tcmalloc/chromium/src/heap-profiler.cc b/third_party/tcmalloc/chromium/src/heap-profiler.cc
index e08b0fcc7cbd233d0b3b73dba2dc7df519f86d16..461b81376260f9af3ea4c44587fc866b09381b1a 100644
--- a/third_party/tcmalloc/chromium/src/heap-profiler.cc
+++ b/third_party/tcmalloc/chromium/src/heap-profiler.cc
@@ -220,6 +220,9 @@ static int64 last_dump_time = 0;      // The time of the last dump
 static HeapProfileTable* heap_profile = NULL;  // the heap profile table
 static DeepHeapProfile* deep_profile = NULL;  // deep memory profiler
 
+// Callback an appplication can use to generate its own "stacks".
+static PseudoStackGenerator pseudo_stack_generator = NULL;
+
 //----------------------------------------------------------------------
 // Profile generation
 //----------------------------------------------------------------------
@@ -238,6 +241,8 @@ static char* DoGetHeapProfileLocked(const char* reason, char* buf, int buflen) {
     (void)stats;   // avoid an unused-variable warning in non-debug mode.
     if (deep_profile) {
       bytes_written = deep_profile->FillOrderedProfile(reason, buf, buflen - 1);
+    } else if (pseudo_stack_generator) {
+      bytes_written = heap_profile->FillPseudoStackProfile(buf, buflen - 1);
     } else {
       bytes_written = heap_profile->FillOrderedProfile(buf, buflen - 1);
     }
@@ -379,6 +384,21 @@ static void RecordAlloc(const void* ptr, size_t bytes, int skip_count) {
   }
 }
 
+// Record an allocation in the profile.
+static void PseudoStackRecordAlloc(const void* ptr,
+                                   size_t bytes,
+                                   int skip_count) {
+  // Take the stack trace outside the critical section.
+  void* stack[HeapProfileTable::kMaxStackDepth];
+  // Generate our pseudo-stack by via callback into the client code.
+  int depth = (*pseudo_stack_generator)(stack);
+  SpinLockHolder l(&heap_lock);
+  if (is_on) {
+    heap_profile->RecordAlloc(ptr, bytes, depth, stack);
+    MaybeDumpProfileLocked();
+  }
+}
+
 // Record a deallocation in the profile.
 static void RecordFree(const void* ptr) {
   SpinLockHolder l(&heap_lock);
@@ -402,6 +422,11 @@ void DeleteHook(const void* ptr) {
   if (ptr != NULL) RecordFree(ptr);
 }
 
+// static
+void PseudoStackNewHook(const void* ptr, size_t size) {
+  if (ptr != NULL) PseudoStackRecordAlloc(ptr, size, 0);
+}
+
 // TODO(jandrews): Re-enable stack tracing
 #ifdef TODO_REENABLE_STACK_TRACING
 static void RawInfoStackDumper(const char* message, void*) {
@@ -534,7 +559,11 @@ extern "C" void HeapProfilerStart(const char* prefix) {
 
   if (FLAGS_only_mmap_profile == false) {
     // Now set the hooks that capture new/delete and malloc/free.
-    RAW_CHECK(MallocHook::AddNewHook(&NewHook), "");
+    if (pseudo_stack_generator) {
+      RAW_CHECK(MallocHook::AddNewHook(&PseudoStackNewHook), "");
+    } else {
+      RAW_CHECK(MallocHook::AddNewHook(&NewHook), "");
+    }
     RAW_CHECK(MallocHook::AddDeleteHook(&DeleteHook), "");
   }
 
@@ -566,7 +595,11 @@ extern "C" void HeapProfilerStop() {
 
   if (FLAGS_only_mmap_profile == false) {
     // Unset our new/delete hooks, checking they were set:
-    RAW_CHECK(MallocHook::RemoveNewHook(&NewHook), "");
+    if (pseudo_stack_generator) {
+      RAW_CHECK(MallocHook::RemoveNewHook(&PseudoStackNewHook), "");
+    } else {
+      RAW_CHECK(MallocHook::RemoveNewHook(&NewHook), "");
+    }
     RAW_CHECK(MallocHook::RemoveDeleteHook(&DeleteHook), "");
   }
   if (FLAGS_mmap_log) {
@@ -638,6 +671,11 @@ extern "C" void HeapProfilerDumpAliveObjects(const char* filename) {
   heap_profile->DumpMarkedObjects(HeapProfileTable::MARK_TWO, filename);
 }
 
+extern "C" void SetPseudoStackGenerator(PseudoStackGenerator callback) {
+  SpinLockHolder l(&heap_lock);
+  pseudo_stack_generator = callback;
+}
+
 //----------------------------------------------------------------------
 // Initialization/finalization code
 //----------------------------------------------------------------------