[memory-infra] Add api to enable heap profiling in MemoryDumpManager

base/trace_event/memory_dump_manager.cc

Index: base/trace_event/memory_dump_manager.cc
diff --git a/base/trace_event/memory_dump_manager.cc b/base/trace_event/memory_dump_manager.cc
index 6aaa3599720d150bc957cee18d965f5d93ecf7e5..d125fe705a2c4677c0c047d11b0f17beaf899afd 100644
--- a/base/trace_event/memory_dump_manager.cc
+++ b/base/trace_event/memory_dump_manager.cc
@@ -160,7 +160,7 @@ MemoryDumpManager::MemoryDumpManager()
     : memory_tracing_enabled_(0),
       tracing_process_id_(kInvalidTracingProcessId),
       dumper_registrations_ignored_for_testing_(false),
-      heap_profiling_enabled_(false) {
+      heap_profiling_state_(HeapProfilingState::DISABLED) {
   g_next_guid.GetNext();  // Make sure that first guid is not zero.
 
   // At this point the command line may not be initialized but we try to
@@ -176,7 +176,8 @@ MemoryDumpManager::~MemoryDumpManager() {
 }
 
 void MemoryDumpManager::EnableHeapProfilingIfNeeded() {
-  if (heap_profiling_enabled_)
+  AutoLock lock(lock_);
+  if (heap_profiling_state_ != HeapProfilingState::DISABLED)
     return;
 
   if (!CommandLine::InitializedForCurrentProcess() ||
@@ -187,14 +188,14 @@ void MemoryDumpManager::EnableHeapProfilingIfNeeded() {
   std::string profiling_mode = CommandLine::ForCurrentProcess()
       ->GetSwitchValueASCII(switches::kEnableHeapProfiling);
   if (profiling_mode == "") {
-    AllocationContextTracker::SetCaptureMode(
+    EnableHeapProfilingLocked(
         AllocationContextTracker::CaptureMode::PSEUDO_STACK);
 #if HAVE_TRACE_STACK_FRAME_POINTERS && \
     (BUILDFLAG(ENABLE_PROFILING) || !defined(NDEBUG))
   } else if (profiling_mode == switches::kEnableHeapProfilingModeNative) {
     // We need frame pointers for native tracing to work, and they are
     // enabled in profiling and debug builds.
-    AllocationContextTracker::SetCaptureMode(
+    EnableHeapProfilingLocked(
         AllocationContextTracker::CaptureMode::NATIVE_STACK);
 #endif
 #if BUILDFLAG(ENABLE_MEMORY_TASK_PROFILER)
@@ -208,10 +209,42 @@ void MemoryDumpManager::EnableHeapProfilingIfNeeded() {
     CHECK(false) << "Invalid mode '" << profiling_mode << "' for "
                << switches::kEnableHeapProfiling << " flag.";
   }
+}
 
-  for (auto mdp : dump_providers_)
-    mdp->dump_provider->OnHeapProfilingEnabled(true);
-  heap_profiling_enabled_ = true;
+void MemoryDumpManager::EnableHeapProfiling(
+    AllocationContextTracker::CaptureMode capture_mode) {
+  CHECK_NE(AllocationContextTracker::CaptureMode::NATIVE_STACK, capture_mode)
+      << "Native mode can be enabled only by command line flag";
+  AutoLock lock(lock_);
+  EnableHeapProfilingLocked(capture_mode);
+}
+
+void MemoryDumpManager::EnableHeapProfilingLocked(
+    AllocationContextTracker::CaptureMode capture_mode) {
+  // Do not enable heap profiling if tracing is enabled, since session_state_
+  // will not be initialized.
+  if (capture_mode != AllocationContextTracker::CaptureMode::DISABLED &&
+      subtle::NoBarrier_Load(&memory_tracing_enabled_)) {
+    return;
+  }
+
+  switch (capture_mode) {
+    case AllocationContextTracker::CaptureMode::PSEUDO_STACK:
+    case AllocationContextTracker::CaptureMode::NATIVE_STACK:
+    case AllocationContextTracker::CaptureMode::BACKGROUND:
+      if (heap_profiling_state_ != HeapProfilingState::DISABLED)
+        return;
+      heap_profiling_state_ = HeapProfilingState::ENABLED;
+      break;
+    case AllocationContextTracker::CaptureMode::DISABLED:
+      heap_profiling_state_ = HeapProfilingState::DISABLED_PERMANENTLY;
+      break;
+  }
+  AllocationContextTracker::SetCaptureMode(capture_mode);
+  for (auto mdp : dump_providers_) {
+    mdp->dump_provider->OnHeapProfilingEnabled(heap_profiling_state_ ==
+                                               HeapProfilingState::ENABLED);
+  }
 }
 
 void MemoryDumpManager::Initialize(
@@ -221,8 +254,8 @@ void MemoryDumpManager::Initialize(
     DCHECK(delegate);
     DCHECK(!delegate_);
     delegate_ = std::move(delegate);
-    EnableHeapProfilingIfNeeded();
   }
+  EnableHeapProfilingIfNeeded();
 
 // Enable the core dump providers.
 #if defined(MALLOC_MEMORY_TRACING_SUPPORTED)
@@ -340,7 +373,7 @@ void MemoryDumpManager::RegisterDumpProviderInternal(
     }
   }
 
-  if (heap_profiling_enabled_)
+  if (heap_profiling_state_ == HeapProfilingState::ENABLED)
     mdp->OnHeapProfilingEnabled(true);
 }
 
@@ -852,7 +885,7 @@ void MemoryDumpManager::OnTraceLogEnabled() {
   session_state->SetAllowedDumpModes(memory_dump_config.allowed_dump_modes);
   session_state->set_heap_profiler_breakdown_threshold_bytes(
       memory_dump_config.heap_profiler_options.breakdown_threshold_bytes);
-  if (heap_profiling_enabled_) {
+  if (heap_profiling_state_ == HeapProfilingState::ENABLED) {
     // If heap profiling is enabled, the stack frame deduplicator and type name
     // deduplicator will be in use. Add a metadata events to write the frames
     // and type IDs.