Implement Shared SingleThreadTaskRunners in the Task Scheduler

base/task_scheduler/scheduler_single_thread_task_runner_manager.cc

Index: base/task_scheduler/scheduler_single_thread_task_runner_manager.cc
diff --git a/base/task_scheduler/scheduler_single_thread_task_runner_manager.cc b/base/task_scheduler/scheduler_single_thread_task_runner_manager.cc
index d96b8cef643dfa211d58df4a44214faa3980276b..043be7e1dee1e59d974c24cb6741d3dd5aeecb1a 100644
--- a/base/task_scheduler/scheduler_single_thread_task_runner_manager.cc
+++ b/base/task_scheduler/scheduler_single_thread_task_runner_manager.cc
@@ -283,6 +283,9 @@ class SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner
     return GetDelegate()->RunsTasksInCurrentSequence();
   }
 
+  // SchedulerSingleThreadTaskRunner:

[fdoray, 2017/05/24 13:25:45]

Update comment. This isn't an override.

[robliao, 2017/05/24 18:28:32]

The headers are not strictly used for overrides, but rather to indicate sections
for different classes. Without this header, you can't tell when
SingleThreadTaskRunner ends.

+  void Start() { worker_->Start(); }
+
  private:
   ~SchedulerSingleThreadTaskRunner() override {
     // Note: This will crash if SchedulerSingleThreadTaskRunnerManager is
@@ -324,6 +327,15 @@ SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunnerManager(
     : task_tracker_(task_tracker), delayed_task_manager_(delayed_task_manager) {
   DCHECK(task_tracker_);
   DCHECK(delayed_task_manager_);
+  static_assert(
+      arraysize(shared_task_runners_) == ENVIRONMENT_COUNT,
+      "The size of |shared_task_runners_| must match ENVIRONMENT_COUNT");
+#if defined(OS_WIN)
+  static_assert(
+      arraysize(shared_com_task_runners_) == arraysize(shared_task_runners_),
+      "The size of |shared_com_task_runners_| must match "
+      "|shared_task_runners_|");
+#endif  // defined(OS_WIN)
 }
 
 SchedulerSingleThreadTaskRunnerManager::
@@ -376,9 +388,10 @@ scoped_refptr<SingleThreadTaskRunner>
 SchedulerSingleThreadTaskRunnerManager::CreateSingleThreadTaskRunnerWithTraits(
     const std::string& name,
     ThreadPriority priority_hint,

[fdoray, 2017/05/24 13:25:45]

Now that this file includes "base/task_scheduler/environment_config.h", it could
get the name and priority hint itself instead of receiving them from the caller.

[robliao, 2017/05/24 18:28:32]

Done.

-    const TaskTraits& traits) {
-  return CreateSingleThreadTaskRunnerWithDelegate<SchedulerWorkerDelegate>(
-      name, priority_hint, traits);
+    const TaskTraits& traits,
+    SingleThreadTaskRunnerThreadMode thread_mode) {
+  return CreateTaskRunnerWithTraitsImpl<SchedulerWorkerDelegate>(
+      name, priority_hint, traits, thread_mode);
 }
 
 #if defined(OS_WIN)
@@ -386,13 +399,55 @@ scoped_refptr<SingleThreadTaskRunner>
 SchedulerSingleThreadTaskRunnerManager::CreateCOMSTATaskRunnerWithTraits(
     const std::string& name,
     ThreadPriority priority_hint,
-    const TaskTraits& traits) {
-  return CreateSingleThreadTaskRunnerWithDelegate<SchedulerWorkerCOMDelegate>(
-      name, priority_hint, traits);
+    const TaskTraits& traits,
+    SingleThreadTaskRunnerThreadMode thread_mode) {
+  return CreateTaskRunnerWithTraitsImpl<SchedulerWorkerCOMDelegate>(
+      name, priority_hint, traits, thread_mode);
 }
 #endif  // defined(OS_WIN)
 
+template <typename DelegateType>
+scoped_refptr<
+    SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner>
+SchedulerSingleThreadTaskRunnerManager::CreateTaskRunnerWithTraitsImpl(
+    const std::string& name,
+    ThreadPriority priority_hint,
+    const TaskTraits& traits,
+    SingleThreadTaskRunnerThreadMode thread_mode) {

[fdoray, 2017/05/24 13:25:45]

I would assert that WithBaseSyncPrimitives() is not used for SHARED
SingleThreadTaskRunners. It could lead to deadlocks (waiting for an event to be
signaled by a task that must run on the same shared thread) or jank.

[robliao, 2017/05/24 18:28:32]

Sounds reasonable to me. Done.

+  // To simply the code, |dedicated_task_runner| is a local only variable that

[fdoray, 2017/05/24 13:25:45]

To *simplify*

[robliao, 2017/05/24 18:28:32]

Done.

+  // allows the code to treat both the DEDICATED and SHARED cases similarly. In
+  // DEDICATED, the scoped_refptr is backed by a local variable and in SHARED,
+  // the scoped_refptr is backed by a member variable.
+  scoped_refptr<SchedulerSingleThreadTaskRunner> dedicated_task_runner;
+  scoped_refptr<SchedulerSingleThreadTaskRunner>& task_runner =
+      thread_mode == SingleThreadTaskRunnerThreadMode::DEDICATED
+          ? dedicated_task_runner
+          : GetSharedTaskRunnerForTraits<DelegateType>(traits);
+  bool new_task_runner = false;
+  bool started;
+  {
+    AutoSchedulerLock auto_lock(lock_);
+    if (!task_runner) {
+      std::string processed_name =
+          thread_mode == SingleThreadTaskRunnerThreadMode::DEDICATED
+              ? name
+              : "Shared" + name;
+      task_runner = CreateSingleThreadTaskRunnerWithDelegate<DelegateType>(
+          processed_name, priority_hint, traits);

[fdoray, 2017/05/24 13:25:45]

Unfortunately, if someone creates a SingleThreadTaskRunner which maps to the
same task environment as an existing SingleThreadTaskRunner but which has
different traits, these new traits won't be honored.

E.g.

tr1 =
CreateSingleThreadTaskRunnerWithTraits({base::TaskShutdownBehavior::BLOCK_SHUTDOWN});
tr2 =
CreateSingleThreadTaskRunnerWithTraits({base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN});
EXPECT_NE(tr1, tr2); // fails

You can reuse SchedulerWorkers, but not SingleThreadTaskRunners.

[robliao, 2017/05/24 18:28:32]

Nice catch. Fixed to reuse SchedulerWorkers.

+      new_task_runner = true;
+    }
+    started = started_;
+  }
+
+  if (new_task_runner && started)
+    task_runner->Start();

[fdoray, 2017/05/24 13:25:45]

What do you think of rewriting this method as:

template <typename DelegateType>
scoped_refptr<
    SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner>
SchedulerSingleThreadTaskRunnerManager::CreateTaskRunnerWithTraitsImpl(
    const std::string& name,
    ThreadPriority priority_hint,
    const TaskTraits& traits,
    SingleThreadTaskRunnerThreadMode thread_mode) {
  // To simply the code, |dedicated_task_runner| is a local only variable that
  // allows the code to treat both the DEDICATED and SHARED cases similarly. In
  // DEDICATED, the scoped_refptr is backed by a local variable and in SHARED,
  // the scoped_refptr is backed by a member variable.
  scoped_refptr<SchedulerSingleThreadTaskRunner> dedicated_task_runner;
  scoped_refptr<SchedulerSingleThreadTaskRunner>& task_runner =
      thread_mode == SingleThreadTaskRunnerThreadMode::DEDICATED
          ? dedicated_task_runner
          : GetSharedTaskRunnerForTraits<DelegateType>(traits);

-- begin changes --

  scoped_refptr<SchedulerWorker> worker;
  bool started;
  {
    AutoSchedulerLock auto_lock(lock_);
    if (!task_runner) {
      std::string processed_name =
          thread_mode == SingleThreadTaskRunnerThreadMode::DEDICATED
              ? name
              : "Shared" + name;
      worker = CreateAndRegisterSchedulerWorker<DelegateType>(
          name, priority_hint);
      task_runner = new SchedulerSingleThreadTaskRunner(this, traits, worker);
    }
    started = started_;
  }
  if (worker && started)
    worker->Start();

-- end changes --

  return task_runner;
}
  
This allows you to delete
SchedulerSingleThreadTaskRunnerManager::CreateSingleThreadTaskRunnerWithDelegate()
and SchedulerSingleThreadTaskRunner::Start() and leads to less lines of code
overall.

[robliao, 2017/05/24 18:28:32]

sgtm. Done.

+
+  return task_runner;
+}
+
 void SchedulerSingleThreadTaskRunnerManager::JoinForTesting() {
+  ReleaseSharedTaskRunners();

[fdoray, 2017/05/24 13:25:45]

Why is this needed? SchedulerWorker::JoinForTesting() should work even if shared
task runners aren't released. Not releasing shared task runners would allows
tasks running during JoinForTesting() to get shared task runners that already
exist without hitting the DCHECK below, which isn't a bad thing.

[robliao, 2017/05/24 18:28:32]

This is now ReleaseSharedSchedulerWorkers() and is needed because the
expectation is that all STTRs are gone by the time
~SchedulerSingleThreadTaskRunnerManager returns. Without this, the workers
remain registered (even though they are joined) and we'll assert at destruction.

+
   decltype(workers_) local_workers;
   {
     AutoSchedulerLock auto_lock(lock_);
@@ -411,7 +466,9 @@ void SchedulerSingleThreadTaskRunnerManager::JoinForTesting() {
 }
 
 template <typename DelegateType>
-scoped_refptr<SingleThreadTaskRunner> SchedulerSingleThreadTaskRunnerManager::
+scoped_refptr<
+    SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner>
+SchedulerSingleThreadTaskRunnerManager::
     CreateSingleThreadTaskRunnerWithDelegate(const std::string& name,
                                              ThreadPriority priority_hint,
                                              const TaskTraits& traits) {
@@ -444,24 +501,31 @@ SchedulerWorker*
 SchedulerSingleThreadTaskRunnerManager::CreateAndRegisterSchedulerWorker(
     const std::string& name,
     ThreadPriority priority_hint) {
-  SchedulerWorker* worker;
-  bool start_worker;
-
-  {
-    AutoSchedulerLock auto_lock(lock_);
-    int id = next_worker_id_++;
-    workers_.emplace_back(make_scoped_refptr(new SchedulerWorker(
-        priority_hint, CreateSchedulerWorkerDelegate<DelegateType>(name, id),
-        task_tracker_)));
-    worker = workers_.back().get();
-    start_worker = started_;
-  }
+  lock_.AssertAcquired();
+  int id = next_worker_id_++;
+  workers_.emplace_back(make_scoped_refptr(new SchedulerWorker(
+      priority_hint, CreateSchedulerWorkerDelegate<DelegateType>(name, id),
+      task_tracker_)));
+  return workers_.back().get();
+}
 
-  if (start_worker)
-    worker->Start();
+template <>
+scoped_refptr<
+    SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner>&
+SchedulerSingleThreadTaskRunnerManager::GetSharedTaskRunnerForTraits<
+    SchedulerWorkerDelegate>(const TaskTraits& traits) {
+  return shared_task_runners_[GetEnvironmentIndexForTraits(traits)];
+}
 
-  return worker;
+#if defined(OS_WIN)
+template <>
+scoped_refptr<
+    SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner>&
+SchedulerSingleThreadTaskRunnerManager::GetSharedTaskRunnerForTraits<
+    SchedulerWorkerCOMDelegate>(const TaskTraits& traits) {
+  return shared_com_task_runners_[GetEnvironmentIndexForTraits(traits)];
 }
+#endif  // defined(OS_WIN)
 
 void SchedulerSingleThreadTaskRunnerManager::UnregisterSchedulerWorker(
     SchedulerWorker* worker) {
@@ -492,5 +556,21 @@ void SchedulerSingleThreadTaskRunnerManager::UnregisterSchedulerWorker(
   worker_to_destroy->Cleanup();
 }
 
+void SchedulerSingleThreadTaskRunnerManager::ReleaseSharedTaskRunners() {
+  decltype(shared_task_runners_) local_shared_task_runners;
+#if defined(OS_WIN)
+  decltype(shared_com_task_runners_) local_shared_com_task_runners;
+#endif
+  {
+    AutoSchedulerLock auto_lock(lock_);
+    for (size_t i = 0; i < arraysize(shared_task_runners_); ++i) {
+      local_shared_task_runners[i] = std::move(shared_task_runners_[i]);
+#if defined(OS_WIN)
+      local_shared_com_task_runners[i] = std::move(shared_com_task_runners_[i]);
+#endif
+    }
+  }
+}
+
 }  // namespace internal
 }  // namespace base