Implement Shared SingleThreadTaskRunners in the Task Scheduler

base/task_scheduler/scheduler_single_thread_task_runner_manager.cc

 // Copyright 2017 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/task_scheduler/scheduler_single_thread_task_runner_manager.h"
 
 #include <algorithm>
 #include <memory>
 #include <string>
 #include <utility>
@@ skipping 265 matching lines @@
                                   OnceClosure closure,
                                   TimeDelta delay) override {
     // Tasks are never nested within the task scheduler.
     return PostDelayedTask(from_here, std::move(closure), delay);
   }
 
   bool RunsTasksInCurrentSequence() const override {
     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
     // incorrectly destroyed first in tests (in production the TaskScheduler and
     // all of its state are intentionally leaked after
     // TaskScheduler::Shutdown(). See ~SchedulerSingleThreadTaskRunnerManager()
     // for more details.
     outer_->UnregisterSchedulerWorker(worker_);
   }
 
@@ skipping 21 matching lines @@
 
   DISALLOW_COPY_AND_ASSIGN(SchedulerSingleThreadTaskRunner);
 };
 
 SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunnerManager(
     TaskTracker* task_tracker,
     DelayedTaskManager* delayed_task_manager)
     : 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::
     ~SchedulerSingleThreadTaskRunnerManager() {
 #if DCHECK_IS_ON()
   size_t workers_unregistered_during_join =
       subtle::NoBarrier_Load(&workers_unregistered_during_join_);
   // Log an ERROR instead of DCHECK'ing as it's often useful to have both the
   // stack trace of this call and the crash stack trace of the upcoming
   // out-of-order ~SchedulerSingleThreadTaskRunner() call to know what to flip.
@@ skipping 31 matching lines @@
   // workers are started as they are created.
   for (scoped_refptr<SchedulerWorker> worker : workers_to_start) {
     worker->Start();
     worker->WakeUp();
   }
 }
 
 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)
 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_);
     local_workers = std::move(workers_);
   }
 
   for (const auto& worker : local_workers)
     worker->JoinForTesting();
 
   {
     AutoSchedulerLock auto_lock(lock_);
     DCHECK(workers_.empty())
         << "New worker(s) unexpectedly registered during join.";
     workers_ = std::move(local_workers);
   }
 }
 
 template <typename DelegateType>
-scoped_refptr<SingleThreadTaskRunner> SchedulerSingleThreadTaskRunnerManager::
+scoped_refptr<
+    SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner>
+SchedulerSingleThreadTaskRunnerManager::
     CreateSingleThreadTaskRunnerWithDelegate(const std::string& name,
                                              ThreadPriority priority_hint,
                                              const TaskTraits& traits) {
   return new SchedulerSingleThreadTaskRunner(
       this, traits,
       CreateAndRegisterSchedulerWorker<DelegateType>(name, priority_hint));
 }
 
 template <>
 std::unique_ptr<SchedulerWorkerDelegate>
@@ skipping 12 matching lines @@
       StringPrintf("TaskSchedulerSingleThreadCOMSTA%s%d", name.c_str(), id),
       task_tracker_);
 }
 #endif  // defined(OS_WIN)
 
 template <typename DelegateType>
 SchedulerWorker*
 SchedulerSingleThreadTaskRunnerManager::CreateAndRegisterSchedulerWorker(
     const std::string& name,
     ThreadPriority priority_hint) {
-  SchedulerWorker* worker;
-  bool start_worker;
+  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();
+}
 
-  {
-    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_;
-  }
+template <>
+scoped_refptr<
+    SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner>&
+SchedulerSingleThreadTaskRunnerManager::GetSharedTaskRunnerForTraits<
+    SchedulerWorkerDelegate>(const TaskTraits& traits) {
+  return shared_task_runners_[GetEnvironmentIndexForTraits(traits)];
+}
 
-  if (start_worker)
-    worker->Start();
-
-  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) {
   // Cleanup uses a SchedulerLock, so call Cleanup() after releasing
   // |lock_|.
   scoped_refptr<SchedulerWorker> worker_to_destroy;
   {
     AutoSchedulerLock auto_lock(lock_);
 
     // We might be joining, so record that a worker was unregistered for
@@ skipping 10 matching lines @@
                      [worker](const scoped_refptr<SchedulerWorker>& candidate) {
                        return candidate.get() == worker;
                      });
     DCHECK(worker_iter != workers_.end());
     worker_to_destroy = std::move(*worker_iter);
     workers_.erase(worker_iter);
   }
   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