Take advantage of multi-threaded MOCK_TIME in HeapProfilerControllerTest

First use case for
https://chromium-review.googlesource.com/c/chromium/src/+/1686776

Bug: 946657
Change-Id: Ic7d6e57c4d2bbf020340690017bbce5af155cfe1
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1687272
Commit-Queue: Gabriel Charette <gab@chromium.org>
Reviewed-by: Alexei Filippov <alph@chromium.org>
Cr-Commit-Position: refs/heads/master@{#676389}

chrome/common/heap_profiler_controller.cc

@@ -13,6 +13,7 @@
 #include "base/sampling_heap_profiler/module_cache.h"
 #include "base/sampling_heap_profiler/sampling_heap_profiler.h"
 #include "base/task/post_task.h"
+#include "base/threading/sequenced_task_runner_handle.h"
 #include "components/metrics/call_stack_profile_builder.h"
 #include "components/metrics/call_stack_profile_metrics_provider.h"
 
@@ -49,34 +50,25 @@ void HeapProfilerController::Start() {
       base::SamplingHeapProfiler::Get()->SetSamplingInterval(sampling_rate);
     base::SamplingHeapProfiler::Get()->Start();
   }
-  ScheduleNextSnapshot(task_runner_ ? std::move(task_runner_)
-                                    : base::CreateTaskRunnerWithTraits(
-                                          {base::TaskPriority::BEST_EFFORT}),
-                       stopped_);
+  ScheduleNextSnapshot(stopped_);
 }
 
 // static
 void HeapProfilerController::ScheduleNextSnapshot(
-    scoped_refptr<base::TaskRunner> task_runner,
     scoped_refptr<StoppedFlag> stopped) {
-  // TODO(https://crbug.com/946657): Remove the task_runner and replace the call
-  // with base::PostDelayedTaskWithTraits once test::ScopedTaskEnvironment
-  // supports mock time in thread pools.
-  task_runner->PostDelayedTask(
-      FROM_HERE,
-      base::BindOnce(&HeapProfilerController::TakeSnapshot,
-                     std::move(task_runner), std::move(stopped)),
+  base::PostDelayedTask(
+      FROM_HERE, {base::TaskPriority::BEST_EFFORT},
+      base::BindOnce(&HeapProfilerController::TakeSnapshot, std::move(stopped)),
       RandomInterval(kHeapCollectionInterval));
 }
 
 // static
 void HeapProfilerController::TakeSnapshot(
-    scoped_refptr<base::TaskRunner> task_runner,
     scoped_refptr<StoppedFlag> stopped) {
   if (stopped->data.IsSet())
     return;
   RetrieveAndSendSnapshot();
-  ScheduleNextSnapshot(std::move(task_runner), std::move(stopped));
+  ScheduleNextSnapshot(std::move(stopped));
 }
 
 // static

chrome/common/heap_profiler_controller.h

@@ -9,10 +9,6 @@
 #include "base/memory/ref_counted.h"
 #include "base/synchronization/atomic_flag.h"
 
-namespace base {
-class TaskRunner;
-}  // namespace base
-
 // HeapProfilerController controls collection of sampled heap allocation
 // snapshots for the current process.
 class HeapProfilerController {
@@ -23,20 +19,13 @@ class HeapProfilerController {
   // Starts periodic heap snapshot collection.
   void Start();
 
-  void SetTaskRunnerForTest(scoped_refptr<base::TaskRunner> task_runner) {
-    task_runner_ = std::move(task_runner);
-  }
-
  private:
   using StoppedFlag = base::RefCountedData<base::AtomicFlag>;
 
-  static void ScheduleNextSnapshot(scoped_refptr<base::TaskRunner> task_runner,
-                                   scoped_refptr<StoppedFlag> stopped);
-  static void TakeSnapshot(scoped_refptr<base::TaskRunner> task_runner,
-                           scoped_refptr<StoppedFlag> stopped);
+  static void ScheduleNextSnapshot(scoped_refptr<StoppedFlag> stopped);
+  static void TakeSnapshot(scoped_refptr<StoppedFlag> stopped);
   static void RetrieveAndSendSnapshot();
 
-  scoped_refptr<base::TaskRunner> task_runner_;
   scoped_refptr<StoppedFlag> stopped_;
 
   DISALLOW_COPY_AND_ASSIGN(HeapProfilerController);

chrome/common/heap_profiler_controller_unittest.cc

@@ -6,7 +6,7 @@
 
 #include "base/sampling_heap_profiler/sampling_heap_profiler.h"
 #include "base/test/bind_test_util.h"
-#include "base/test/test_mock_time_task_runner.h"
+#include "base/test/scoped_task_environment.h"
 #include "build/build_config.h"
 #include "components/metrics/call_stack_profile_builder.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -19,31 +19,31 @@
 #define MAYBE_EmptyProfileIsNotEmitted EmptyProfileIsNotEmitted
 #endif
 
-TEST(HeapProfilerControllerTest, MAYBE_EmptyProfileIsNotEmitted) {
-  auto task_runner = base::MakeRefCounted<base::TestMockTimeTaskRunner>();
-  base::TestMockTimeTaskRunner::ScopedContext scoped_context(task_runner.get());
+class HeapProfilerControllerTest : public testing::Test {
+ protected:
+  base::test::ScopedTaskEnvironment scoped_task_environment{
+      base::test::ScopedTaskEnvironment::MainThreadType::MOCK_TIME,
+      base::test::ScopedTaskEnvironment::NowSource::MAIN_THREAD_MOCK_TIME};
+};
 
+TEST_F(HeapProfilerControllerTest, MAYBE_EmptyProfileIsNotEmitted) {
   HeapProfilerController controller;
   metrics::CallStackProfileBuilder::SetBrowserProcessReceiverCallback(
       base::BindLambdaForTesting(
           [](base::TimeTicks time, metrics::SampledProfile profile) {
             ADD_FAILURE();
           }));
-  controller.SetTaskRunnerForTest(task_runner);
   controller.Start();
 
-  task_runner->FastForwardBy(base::TimeDelta::FromDays(365));
+  scoped_task_environment.FastForwardBy(base::TimeDelta::FromDays(365));
 }
 
 // Sampling profiler is not capable of unwinding stack on Android under tests.
 #if !defined(OS_ANDROID)
-TEST(HeapProfilerControllerTest, ProfileCollectionsScheduler) {
+TEST_F(HeapProfilerControllerTest, ProfileCollectionsScheduler) {
   constexpr size_t kAllocationSize = 42 * 1024;
   constexpr int kSnapshotsToCollect = 3;
 
-  auto task_runner = base::MakeRefCounted<base::TestMockTimeTaskRunner>();
-  base::TestMockTimeTaskRunner::ScopedContext scoped_context(task_runner.get());
-
   auto controller = std::make_unique<HeapProfilerController>();
   int profile_count = 0;
 
@@ -75,7 +75,6 @@ TEST(HeapProfilerControllerTest, ProfileCollectionsScheduler) {
 
   metrics::CallStackProfileBuilder::SetBrowserProcessReceiverCallback(
       base::BindLambdaForTesting(check_profile));
-  controller->SetTaskRunnerForTest(task_runner);
   controller->Start();
 
   auto* sampler = base::PoissonAllocationSampler::Get();
@@ -85,7 +84,7 @@ TEST(HeapProfilerControllerTest, ProfileCollectionsScheduler) {
   sampler->RecordAlloc(reinterpret_cast<void*>(0x7331), kAllocationSize,
                        base::PoissonAllocationSampler::kMalloc, nullptr);
 
-  task_runner->FastForwardUntilNoTasksRemain();
+  scoped_task_environment.FastForwardUntilNoTasksRemain();
   EXPECT_LE(kSnapshotsToCollect, profile_count);
 }
 #endif