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

base/trace_event/memory_dump_manager_unittest.cc

Index: base/trace_event/memory_dump_manager_unittest.cc
diff --git a/base/trace_event/memory_dump_manager_unittest.cc b/base/trace_event/memory_dump_manager_unittest.cc
index 3a2992b6f92d687d72d8c0902501ca041ad99389..0b135af0133a3c898a3e49378538bede480a1216 100644
--- a/base/trace_event/memory_dump_manager_unittest.cc
+++ b/base/trace_event/memory_dump_manager_unittest.cc
@@ -9,7 +9,9 @@
 #include <memory>
 #include <vector>
 
+#include "base/base_switches.h"
 #include "base/bind_helpers.h"
+#include "base/command_line.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/ref_counted_memory.h"
 #include "base/message_loop/message_loop.h"
@@ -144,6 +146,7 @@ class MockMemoryDumpProvider : public MemoryDumpProvider {
   MOCK_METHOD0(Destructor, void());
   MOCK_METHOD2(OnMemoryDump,
                bool(const MemoryDumpArgs& args, ProcessMemoryDump* pmd));
+  MOCK_METHOD1(OnHeapProfilingEnabled, void(bool enabled));
   MOCK_METHOD1(PollFastMemoryTotal, void(uint64_t* memory_total));
   MOCK_METHOD0(SuspendFastMemoryPolling, void());
 
@@ -1292,5 +1295,75 @@ TEST_F(MemoryDumpManagerTest, TestBlacklistedUnsafeUnregistration) {
   thread.Stop();
 }
 
+TEST_F(MemoryDumpManagerTest, EnableHeapProfiling_PSEUDO_STACK) {
+  InitializeMemoryDumpManager(false /* is_coordinator */);
+  MockMemoryDumpProvider mdp1;
+  RegisterDumpProvider(&mdp1, nullptr);
+  testing::InSequence sequence;
+  EXPECT_CALL(mdp1, OnHeapProfilingEnabled(true)).Times(1);
+  EXPECT_CALL(mdp1, OnHeapProfilingEnabled(false)).Times(1);
+
+  EnableTracingWithLegacyCategories(MemoryDumpManager::kTraceCategory);
+  mdm_->EnableHeapProfiling(
+      AllocationContextTracker::CaptureMode::PSEUDO_STACK);
+  ASSERT_EQ(AllocationContextTracker::CaptureMode::DISABLED,
+            AllocationContextTracker::capture_mode());
+  DisableTracing();
+
+  mdm_->EnableHeapProfiling(
+      AllocationContextTracker::CaptureMode::PSEUDO_STACK);
+  ASSERT_EQ(AllocationContextTracker::CaptureMode::PSEUDO_STACK,
+            AllocationContextTracker::capture_mode());
+  mdm_->EnableHeapProfiling(AllocationContextTracker::CaptureMode::DISABLED);
+  ASSERT_EQ(AllocationContextTracker::CaptureMode::DISABLED,
+            AllocationContextTracker::capture_mode());
+  mdm_->EnableHeapProfiling(
+      AllocationContextTracker::CaptureMode::PSEUDO_STACK);
+  ASSERT_EQ(AllocationContextTracker::CaptureMode::DISABLED,
+            AllocationContextTracker::capture_mode());
+}
+
+TEST_F(MemoryDumpManagerTest, EnableHeapProfiling_BACKGROUND) {
+  InitializeMemoryDumpManager(true /* is_coordinator */);
+  MockMemoryDumpProvider mdp1;
+  RegisterDumpProvider(&mdp1, nullptr);
+  testing::InSequence sequence;
+  EXPECT_CALL(mdp1, OnHeapProfilingEnabled(true)).Times(1);
+  EXPECT_CALL(mdp1, OnHeapProfilingEnabled(false)).Times(1);
+
+  mdm_->EnableHeapProfiling(AllocationContextTracker::CaptureMode::BACKGROUND);
+  ASSERT_EQ(AllocationContextTracker::CaptureMode::BACKGROUND,
+            AllocationContextTracker::capture_mode());
+  mdm_->EnableHeapProfiling(AllocationContextTracker::CaptureMode::DISABLED);
+  ASSERT_EQ(AllocationContextTracker::CaptureMode::DISABLED,
+            AllocationContextTracker::capture_mode());
+  mdm_->EnableHeapProfiling(
+      AllocationContextTracker::CaptureMode::PSEUDO_STACK);
+  ASSERT_EQ(AllocationContextTracker::CaptureMode::DISABLED,
+            AllocationContextTracker::capture_mode());
+}
+
+TEST_F(MemoryDumpManagerTest, EnableHeapProfilingIfNeeded) {
+  InitializeMemoryDumpManager(false /* is_coordinator */);
+  MockMemoryDumpProvider mdp1;
+  RegisterDumpProvider(&mdp1, nullptr);
+  testing::InSequence sequence;
+  EXPECT_CALL(mdp1, OnHeapProfilingEnabled(true)).Times(1);
+  EXPECT_CALL(mdp1, OnHeapProfilingEnabled(false)).Times(1);
+
+  CommandLine* cmdline = CommandLine::ForCurrentProcess();
+  cmdline->AppendSwitchASCII(switches::kEnableHeapProfiling, "");
+  mdm_->EnableHeapProfilingIfNeeded();
+  ASSERT_EQ(AllocationContextTracker::CaptureMode::PSEUDO_STACK,
+            AllocationContextTracker::capture_mode());
+  mdm_->EnableHeapProfiling(AllocationContextTracker::CaptureMode::DISABLED);
+  ASSERT_EQ(AllocationContextTracker::CaptureMode::DISABLED,
+            AllocationContextTracker::capture_mode());
+  mdm_->EnableHeapProfiling(
+      AllocationContextTracker::CaptureMode::PSEUDO_STACK);
+  ASSERT_EQ(AllocationContextTracker::CaptureMode::DISABLED,
+            AllocationContextTracker::capture_mode());
+}
+
 }  // namespace trace_event
 }  // namespace base