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 228 matching lines @@
 }  // namespace
 
 class SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner
     : public SingleThreadTaskRunner {
  public:
   // Constructs a SchedulerSingleThreadTaskRunner that indirectly controls the
   // lifetime of a dedicated |worker| for |traits|.
   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_);
   }
 
   // SingleThreadTaskRunner:
   bool PostDelayedTask(const tracked_objects::Location& from_here,
                        OnceClosure closure,
                        TimeDelta delay) override {
     auto task = MakeUnique<Task>(from_here, std::move(closure), traits_, delay);
     task->single_thread_task_runner_ref = this;
@@ skipping 17 matching lines @@
     // Tasks are never nested within the task scheduler.
     return PostDelayedTask(from_here, std::move(closure), delay);
   }
 
   bool RunsTasksInCurrentSequence() const override {
     return GetDelegate()->RunsTasksInCurrentSequence();
   }
 
  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) {
     scoped_refptr<Sequence> sequence = GetDelegate()->sequence();
     // If |sequence| is null, then the thread is effectively gone (either
     // shutdown or joined).
     if (!sequence)
       return;
 
     const bool sequence_was_empty = sequence->PushTask(std::move(task));
     if (sequence_was_empty) {
       GetDelegate()->ReEnqueueSequence(std::move(sequence));
       worker_->WakeUp();
     }
   }
 
   SchedulerWorkerDelegate* GetDelegate() const {
     return static_cast<SchedulerWorkerDelegate*>(worker_->delegate());
   }
 
   SchedulerSingleThreadTaskRunnerManager* const outer_;
   const TaskTraits traits_;
   SchedulerWorker* const worker_;
+  const SingleThreadTaskRunnerThreadMode thread_mode_;
 
   DISALLOW_COPY_AND_ASSIGN(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::
     ~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,
-    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_);
     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::
-    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<
     SchedulerWorkerDelegate>(const std::string& name, int id) {
   return MakeUnique<SchedulerWorkerDelegate>(
       StringPrintf("TaskSchedulerSingleThread%s%d", name.c_str(), id));
 }
 
 #if defined(OS_WIN)
 template <>
 std::unique_ptr<SchedulerWorkerDelegate>
 SchedulerSingleThreadTaskRunnerManager::CreateSchedulerWorkerDelegate<
     SchedulerWorkerCOMDelegate>(const std::string& name, int id) {
   return MakeUnique<SchedulerWorkerCOMDelegate>(
       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 <>
+SchedulerWorker*&
+SchedulerSingleThreadTaskRunnerManager::GetSharedSchedulerWorkerForTraits<
+    SchedulerWorkerDelegate>(const TaskTraits& traits) {
+  return shared_scheduler_workers_[GetEnvironmentIndexForTraits(traits)];
+}
 
-  if (start_worker)
-    worker->Start();
-
-  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) {
   // 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::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