Run ScopedTaskScheduler tasks synchronously.

base/test/scoped_task_scheduler.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/test/scoped_task_scheduler.h"
 
+#include <memory>
+#include <utility>
 #include <vector>
 
 #include "base/bind.h"
-#include "base/task_scheduler/scheduler_worker_pool_params.h"
+#include "base/logging.h"
+#include "base/memory/ptr_util.h"
+#include "base/memory/ref_counted.h"
+#include "base/message_loop/message_loop.h"
+#include "base/metrics/histogram_base.h"
+#include "base/run_loop.h"
+#include "base/sequence_token.h"
+#include "base/sequenced_task_runner.h"
+#include "base/single_thread_task_runner.h"
+#include "base/task_runner.h"
+#include "base/task_scheduler/task.h"
 #include "base/task_scheduler/task_scheduler.h"
-#include "base/threading/platform_thread.h"
-#include "base/time/time.h"
+#include "base/task_scheduler/task_tracker.h"
+#include "base/threading/thread_task_runner_handle.h"
 
 namespace base {
 namespace test {
 
+namespace {
+
+enum class ExecutionMode { PARALLEL, SEQUENCED, SINGLE_THREADED };
+
+class TestTaskScheduler : public TaskScheduler {
+ public:
+  TestTaskScheduler();
+  ~TestTaskScheduler() override;
+
+  // TaskScheduler:
+  void PostTaskWithTraits(const tracked_objects::Location& from_here,
+                          const TaskTraits& traits,
+                          const Closure& task) override;
+  scoped_refptr<TaskRunner> CreateTaskRunnerWithTraits(
+      const TaskTraits& traits) override;
+  scoped_refptr<SequencedTaskRunner> CreateSequencedTaskRunnerWithTraits(
+      const TaskTraits& traits) override;
+  scoped_refptr<SingleThreadTaskRunner> CreateSingleThreadTaskRunnerWithTraits(
+      const TaskTraits& traits) override;
+  std::vector<const HistogramBase*> GetHistograms() const override;
+  void Shutdown() override;
+  void FlushForTesting() override;
+
+  // Posts |task| to this TaskScheduler with |sequence_token|. Returns true on
+  // success.
+  bool PostTask(std::unique_ptr<internal::Task> task,
+                const SequenceToken& sequence_token);
+
+  // Runs |task| with |sequence_token| using this TaskScheduler's TaskTracker.
+  void RunTask(std::unique_ptr<internal::Task> task,
+               const SequenceToken& sequence_token);
+
+  // Returns true if this TaskScheduler runs its tasks on the current thread.
+  bool RunsTasksOnCurrentThread() const;
+
+ private:
+  // The MessageLoop created by this TaskScheduler. Nullptr if there was already
+  // a MessageLoop on the thread when this TaskScheduler was created.
+  const std::unique_ptr<MessageLoop> message_loop_ =
+      MessageLoop::current() ? nullptr : MakeUnique<MessageLoop>();

[gab, 2016/12/09 20:04:17]

Not a fan of MessageLoop::current() for all the reasons we're trying to minimize
its usage (e.g. unclear where the global is coming from).

I'd rather have an alternate constructor that takes one explicitly to highlight
when reading that it's sharing an existing MessageLoop (just like base::Thread's
constructor).

I guess that would mean :

// |message_loop_owned_| will be non-null if this TestTaskScheduler owns the
// MessageLoop (wasn't handed one). |message_loop_| will always be set and
// should be used by the class.
std::unique_ptr<MessageLoop> message_loop_owned_;
MessageLoop* message_loop_;

as members though.. but I'm okay paying this cost in the anonymous namespace to
have a more readable API, WDYT?

[fdoray, 2016/12/12 18:18:50]

Done.

+
+  // The SingleThreadTaskRunner associated with |message_loop_|.
+  const scoped_refptr<SingleThreadTaskRunner> message_loop_task_runner_ =
+      ThreadTaskRunnerHandle::Get();
+
+  // Handles shutdown behaviors and sets up the environment to run a task.
+  internal::TaskTracker task_tracker_;
+
+  DISALLOW_COPY_AND_ASSIGN(TestTaskScheduler);
+};
+
+class TestTaskSchedulerTaskRunner : public SingleThreadTaskRunner {
+ public:
+  TestTaskSchedulerTaskRunner(TestTaskScheduler* task_scheduler,
+                              ExecutionMode execution_mode,
+                              TaskTraits traits);
+
+  // SingleThreadTaskRunner:
+  bool PostDelayedTask(const tracked_objects::Location& from_here,
+                       const Closure& closure,
+                       TimeDelta delay) override;
+  bool PostNonNestableDelayedTask(const tracked_objects::Location& from_here,
+                                  const Closure& closure,
+                                  TimeDelta delay) override;
+  bool RunsTasksOnCurrentThread() const override;
+
+ private:
+  ~TestTaskSchedulerTaskRunner() override;
+
+  TestTaskScheduler* const task_scheduler_;
+  const ExecutionMode execution_mode_;
+  const SequenceToken sequence_token_;
+  const TaskTraits traits_;
+
+  DISALLOW_COPY_AND_ASSIGN(TestTaskSchedulerTaskRunner);
+};
+
+TestTaskScheduler::TestTaskScheduler() = default;
+
+TestTaskScheduler::~TestTaskScheduler() {
+  // Prevent the RunUntilIdle() call below from running SKIP_ON_SHUTDOWN and
+  // CONTINUE_ON_SHUTDOWN tasks.
+  task_tracker_.SetHasShutdownStartedForTesting();
+
+  // Run pending BLOCK_SHUTDOWN tasks.
+  RunLoop().RunUntilIdle();
+}
+
+void TestTaskScheduler::PostTaskWithTraits(
+    const tracked_objects::Location& from_here,
+    const TaskTraits& traits,
+    const Closure& task) {
+  CreateTaskRunnerWithTraits(traits)->PostTask(from_here, task);
+}
+
+scoped_refptr<TaskRunner> TestTaskScheduler::CreateTaskRunnerWithTraits(
+    const TaskTraits& traits) {
+  return make_scoped_refptr(
+      new TestTaskSchedulerTaskRunner(this, ExecutionMode::PARALLEL, traits));
+}
+
+scoped_refptr<SequencedTaskRunner>
+TestTaskScheduler::CreateSequencedTaskRunnerWithTraits(
+    const TaskTraits& traits) {
+  return make_scoped_refptr(
+      new TestTaskSchedulerTaskRunner(this, ExecutionMode::SEQUENCED, traits));
+}
+
+scoped_refptr<SingleThreadTaskRunner>
+TestTaskScheduler::CreateSingleThreadTaskRunnerWithTraits(
+    const TaskTraits& traits) {
+  return make_scoped_refptr(new TestTaskSchedulerTaskRunner(
+      this, ExecutionMode::SINGLE_THREADED, traits));
+}
+
+std::vector<const HistogramBase*> TestTaskScheduler::GetHistograms() const {
+  NOTREACHED();
+  return std::vector<const HistogramBase*>();
+}
+
+void TestTaskScheduler::Shutdown() {
+  NOTREACHED();
+}
+
+void TestTaskScheduler::FlushForTesting() {
+  NOTREACHED();
+}
+
+bool TestTaskScheduler::PostTask(std::unique_ptr<internal::Task> task,
+                                 const SequenceToken& sequence_token) {
+  DCHECK(task);
+  if (!task_tracker_.WillPostTask(task.get()))
+    return false;
+  internal::Task* const task_ptr = task.get();
+  return message_loop_task_runner_->PostDelayedTask(
+      task_ptr->posted_from, Bind(&TestTaskScheduler::RunTask, Unretained(this),
+                                  Passed(std::move(task)), sequence_token),
+      task_ptr->delay);
+}
+
+void TestTaskScheduler::RunTask(std::unique_ptr<internal::Task> task,
+                                const SequenceToken& sequence_token) {
+  // Clear the MessageLoop TaskRunner to allow TaskTracker to register a
+  // Thread/SequencedTaskRunnerHandle as appropriate.
+  MessageLoop::current()->SetTaskRunner(nullptr);
+
+  // Run the task.
+  task_tracker_.RunTask(std::move(task), sequence_token.IsValid()
+                                             ? sequence_token
+                                             : SequenceToken::Create());
+
+  // Restore the MessageLoop TaskRunner.
+  MessageLoop::current()->SetTaskRunner(message_loop_task_runner_);
+}
+
+bool TestTaskScheduler::RunsTasksOnCurrentThread() const {
+  return message_loop_task_runner_->RunsTasksOnCurrentThread();
+}
+
+TestTaskSchedulerTaskRunner::TestTaskSchedulerTaskRunner(
+    TestTaskScheduler* task_scheduler,
+    ExecutionMode execution_mode,
+    TaskTraits traits)
+    : task_scheduler_(task_scheduler),
+      execution_mode_(execution_mode),
+      sequence_token_(execution_mode == ExecutionMode::PARALLEL
+                          ? SequenceToken()
+                          : SequenceToken::Create()),
+      traits_(traits) {}
+
+bool TestTaskSchedulerTaskRunner::PostDelayedTask(
+    const tracked_objects::Location& from_here,
+    const Closure& closure,
+    TimeDelta delay) {
+  auto task = MakeUnique<internal::Task>(from_here, closure, traits_, delay);
+  if (execution_mode_ == ExecutionMode::SEQUENCED)
+    task->sequenced_task_runner_ref = make_scoped_refptr(this);
+  else if (execution_mode_ == ExecutionMode::SINGLE_THREADED)
+    task->single_thread_task_runner_ref = make_scoped_refptr(this);
+  return task_scheduler_->PostTask(std::move(task), sequence_token_);
+}
+
+bool TestTaskSchedulerTaskRunner::PostNonNestableDelayedTask(
+    const tracked_objects::Location& from_here,
+    const Closure& closure,
+    TimeDelta delay) {
+  // Tasks are never nested within the task scheduler.
+  return PostDelayedTask(from_here, closure, delay);
+}
+
+bool TestTaskSchedulerTaskRunner::RunsTasksOnCurrentThread() const {
+  if (execution_mode_ == ExecutionMode::PARALLEL)
+    return task_scheduler_->RunsTasksOnCurrentThread();
+  return sequence_token_ == SequenceToken::GetForCurrentThread();
+}
+
+TestTaskSchedulerTaskRunner::~TestTaskSchedulerTaskRunner() = default;
+
+}  // namespace
+
 ScopedTaskScheduler::ScopedTaskScheduler() {
   DCHECK(!TaskScheduler::GetInstance());
-
-  // Create a TaskScheduler with a single thread to make tests deterministic.
-  constexpr int kMaxThreads = 1;
-  std::vector<SchedulerWorkerPoolParams> worker_pool_params_vector;
-  worker_pool_params_vector.emplace_back(
-      "Simple", ThreadPriority::NORMAL,
-      SchedulerWorkerPoolParams::StandbyThreadPolicy::LAZY, kMaxThreads,
-      TimeDelta::Max());
-  TaskScheduler::CreateAndSetDefaultTaskScheduler(
-      worker_pool_params_vector,
-      Bind([](const TaskTraits&) -> size_t { return 0; }));
+  TaskScheduler::SetInstance(MakeUnique<TestTaskScheduler>());
   task_scheduler_ = TaskScheduler::GetInstance();
 }
 
 ScopedTaskScheduler::~ScopedTaskScheduler() {
+  DCHECK(thread_checker_.CalledOnValidThread());
   DCHECK_EQ(task_scheduler_, TaskScheduler::GetInstance());
-  TaskScheduler::GetInstance()->Shutdown();
-  TaskScheduler::GetInstance()->JoinForTesting();
   TaskScheduler::SetInstance(nullptr);
 }
 
 }  // namespace test
 }  // namespace base