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..3f45500e8a6edd714084586a32d6b404e5b33de5 100644
--- a/base/task_scheduler/scheduler_single_thread_task_runner_manager.cc
+++ b/base/task_scheduler/scheduler_single_thread_task_runner_manager.cc
@@ -246,8 +246,12 @@ class SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner
   SchedulerSingleThreadTaskRunner(
       SchedulerSingleThreadTaskRunnerManager* const outer,
       const TaskTraits& traits,
-      SchedulerWorker* worker)
-      : outer_(outer), traits_(traits), worker_(worker) {
+      SchedulerWorker* worker,
+      SingleThreadTaskRunnerThreadMode thread_mode)
+      : outer_(outer),
+        traits_(traits),
+        worker_(worker),
+        thread_mode_(thread_mode) {
     DCHECK(outer_);
     DCHECK(worker_);
   }
@@ -285,12 +289,16 @@ class SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner
 
  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_);
+    // Only unregister if this is a DEDICATED SingleThreadTaskRunner. SHARED
+    // task runner SchedulerWorkers are managed separately as they are reused.
+    if (thread_mode_ == SingleThreadTaskRunnerThreadMode::DEDICATED) {
+      // 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_);
+    }
   }
 
   void PostTaskNow(std::unique_ptr<Task> task) {
@@ -314,6 +322,7 @@ class SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner
   SchedulerSingleThreadTaskRunnerManager* const outer_;
   const TaskTraits traits_;
   SchedulerWorker* const worker_;
+  const SingleThreadTaskRunnerThreadMode thread_mode_;
 
   DISALLOW_COPY_AND_ASSIGN(SchedulerSingleThreadTaskRunner);
 };
@@ -321,9 +330,25 @@ class SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner
 SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunnerManager(
     TaskTracker* task_tracker,
     DelayedTaskManager* delayed_task_manager)
-    : task_tracker_(task_tracker), delayed_task_manager_(delayed_task_manager) {
+    : task_tracker_(task_tracker),
+      delayed_task_manager_(delayed_task_manager),
+      shared_scheduler_workers_ {}
+#if defined(OS_WIN)
+      ,
+      shared_com_scheduler_workers_ {}

[gab, 2017/05/25 18:47:00]

Use member initialization for these two as mentioned in .h, also avoids ifdefs
in initializer list :)

[robliao, 2017/05/30 20:01:06]

Done.

+#endif  // defined(OS_WIN)
+{
   DCHECK(task_tracker_);
   DCHECK(delayed_task_manager_);
+  static_assert(
+      arraysize(shared_scheduler_workers_) == ENVIRONMENT_COUNT,

[fdoray, 2017/05/25 17:02:59]

The static_assert is redundant if this is defined as
SchedulerWorker* shared_scheduler_workers_[ENVIRONMENT_COUNT];
in the header file.

[gab, 2017/05/25 18:47:00]

Agreed and with the initialization moving to member initialization these asserts
in the constructor would feel slightly out of place. I think it's safe to drop
them.

[robliao, 2017/05/30 20:01:06]

Done. Keeping the one below though since that certainly needs to remain true and
isn't trivially true by construction.

+      "The size of |shared_scheduler_workers_| must match ENVIRONMENT_COUNT");
+#if defined(OS_WIN)
+  static_assert(arraysize(shared_com_scheduler_workers_) ==
+                    arraysize(shared_scheduler_workers_),
+                "The size of |shared_com_scheduler_workers_| must match "
+                "|shared_scheduler_workers_|");
+#endif  // defined(OS_WIN)
 }
 
 SchedulerSingleThreadTaskRunnerManager::
@@ -375,24 +400,72 @@ void SchedulerSingleThreadTaskRunnerManager::Start() {
 scoped_refptr<SingleThreadTaskRunner>
 SchedulerSingleThreadTaskRunnerManager::CreateSingleThreadTaskRunnerWithTraits(
     const std::string& name,
-    ThreadPriority priority_hint,
-    const TaskTraits& traits) {
-  return CreateSingleThreadTaskRunnerWithDelegate<SchedulerWorkerDelegate>(
-      name, priority_hint, traits);
+    const TaskTraits& traits,
+    SingleThreadTaskRunnerThreadMode thread_mode) {
+  return CreateTaskRunnerWithTraitsImpl<SchedulerWorkerDelegate>(name, 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, traits, thread_mode);
 }
 #endif  // defined(OS_WIN)
 
+template <typename DelegateType>
+scoped_refptr<
+    SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner>
+SchedulerSingleThreadTaskRunnerManager::CreateTaskRunnerWithTraitsImpl(
+    const std::string& name,
+    const TaskTraits& traits,
+    SingleThreadTaskRunnerThreadMode thread_mode) {
+  DCHECK(thread_mode != SingleThreadTaskRunnerThreadMode::SHARED ||
+         !traits.with_base_sync_primitives())
+      << "Using WithBaseSyncPrimitives() on a shared SingleThreadTaskRunner "
+         "may cause deadlocks. Either reevaluate your usage pattern or use "

[gab, 2017/05/25 18:47:00]

"... your usage patterns (e.g. use SequencedTaskRunner) or ..."

[robliao, 2017/05/30 20:01:06]

added the e.g. The singular usage is appropriate here, however.

+         "SingleThreadTaskRunnerThreadMode::DEDICATED.";
+  // To simplify the code, |dedicated_worker| is a local only variable that
+  // allows the code to treat both the DEDICATED and SHARED cases similarly for
+  // SingleThreadTaskRunnerThreadMode. In DEDICATED, the scoped_refptr is backed
+  // by a local variable and in SHARED, the scoped_refptr is backed by a member
+  // variable.
+  SchedulerWorker* dedicated_worker = nullptr;
+  SchedulerWorker*& worker =
+      thread_mode == SingleThreadTaskRunnerThreadMode::DEDICATED
+          ? dedicated_worker
+          : GetSharedSchedulerWorkerForTraits<DelegateType>(traits);
+  bool new_worker = false;
+  bool started;
+  {
+    AutoSchedulerLock auto_lock(lock_);
+    if (!worker) {
+      const auto& environment_params =
+          kEnvironmentParams[GetEnvironmentIndexForTraits(traits)];
+      std::string processed_name =
+          thread_mode == SingleThreadTaskRunnerThreadMode::DEDICATED
+              ? name + environment_params.name_suffix
+              : "Shared" + name + environment_params.name_suffix;
+      worker = CreateAndRegisterSchedulerWorker<DelegateType>(
+          processed_name, environment_params.priority_hint);
+      new_worker = true;
+    }
+    started = started_;
+  }
+
+  if (new_worker && started)
+    worker->Start();
+
+  return new SchedulerSingleThreadTaskRunner(this, traits, worker, thread_mode);

[fdoray, 2017/05/25 17:02:59]

MakeRefCounted<SchedulerSingleThreadTaskRunner>

[robliao, 2017/05/30 20:01:06]

Done.

+}
+
 void SchedulerSingleThreadTaskRunnerManager::JoinForTesting() {
+  ReleaseSharedSchedulerWorkers();
+
   decltype(workers_) local_workers;
   {
     AutoSchedulerLock auto_lock(lock_);
@@ -410,16 +483,6 @@ void SchedulerSingleThreadTaskRunnerManager::JoinForTesting() {
   }
 }
 
-template <typename DelegateType>
-scoped_refptr<SingleThreadTaskRunner> 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>
 SchedulerSingleThreadTaskRunnerManager::CreateSchedulerWorkerDelegate<
@@ -444,24 +507,29 @@ 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 <>
+SchedulerWorker*&
+SchedulerSingleThreadTaskRunnerManager::GetSharedSchedulerWorkerForTraits<
+    SchedulerWorkerDelegate>(const TaskTraits& traits) {
+  return shared_scheduler_workers_[GetEnvironmentIndexForTraits(traits)];
+}
 
-  return worker;
+#if defined(OS_WIN)
+template <>
+SchedulerWorker*&
+SchedulerSingleThreadTaskRunnerManager::GetSharedSchedulerWorkerForTraits<
+    SchedulerWorkerCOMDelegate>(const TaskTraits& traits) {
+  return shared_com_scheduler_workers_[GetEnvironmentIndexForTraits(traits)];
 }
+#endif  // defined(OS_WIN)
 
 void SchedulerSingleThreadTaskRunnerManager::UnregisterSchedulerWorker(
     SchedulerWorker* worker) {
@@ -492,5 +560,30 @@ void SchedulerSingleThreadTaskRunnerManager::UnregisterSchedulerWorker(
   worker_to_destroy->Cleanup();
 }
 
+void SchedulerSingleThreadTaskRunnerManager::ReleaseSharedSchedulerWorkers() {
+  decltype(shared_scheduler_workers_) local_shared_scheduler_workers;
+#if defined(OS_WIN)
+  decltype(shared_com_scheduler_workers_) local_shared_com_scheduler_workers;
+#endif
+  {
+    AutoSchedulerLock auto_lock(lock_);
+    for (size_t i = 0; i < arraysize(shared_scheduler_workers_); ++i) {
+      local_shared_scheduler_workers[i] = shared_scheduler_workers_[i];
+#if defined(OS_WIN)
+      local_shared_com_scheduler_workers[i] = shared_com_scheduler_workers_[i];

[gab, 2017/05/25 18:47:00]

Why do we need a local list? Doesn't seem like anything touches the shared
pointers during UnregisterSchedulerWorker?

Also, shall we set the shared pointers to nullptr when done for good measure?

[robliao, 2017/05/30 20:01:06]

The local list is needed to avoid unregistration within the lock. 

Setting nullptr seemed unnecessary since we make no guarantees after
JoinForTesting. 
I'll go ahead and set them here to root a discussion.

+#endif
+    }
+  }
+
+  for (size_t i = 0; i < arraysize(local_shared_scheduler_workers); ++i) {
+    if (local_shared_scheduler_workers[i])
+      UnregisterSchedulerWorker(local_shared_scheduler_workers[i]);
+#if defined(OS_WIN)
+    if (local_shared_com_scheduler_workers[i])
+      UnregisterSchedulerWorker(local_shared_com_scheduler_workers[i]);
+#endif
+  }
+}
+
 }  // namespace internal
 }  // namespace base