TaskScheduler: Make the worker pools of TaskSchedulerImpl configurable

base/task_scheduler/task_scheduler_impl.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/task_scheduler/task_scheduler_impl.h"
 
 #include <utility>
 
 #include "base/bind.h"
 #include "base/bind_helpers.h"
 #include "base/memory/ptr_util.h"
 #include "base/task_scheduler/scheduler_service_thread.h"
-#include "base/task_scheduler/scheduler_worker_pool_impl.h"
 #include "base/task_scheduler/sequence_sort_key.h"
 #include "base/task_scheduler/task.h"
 #include "base/time/time.h"
 
 namespace base {
 namespace internal {
 
 // static
-std::unique_ptr<TaskSchedulerImpl> TaskSchedulerImpl::Create() {
-  std::unique_ptr<TaskSchedulerImpl> scheduler(new TaskSchedulerImpl);
-  scheduler->Initialize();
+std::unique_ptr<TaskSchedulerImpl> TaskSchedulerImpl::Create(
+    const std::vector<WorkerPoolCreationArgs>& worker_pools,
+    const WorkerPoolIndexForTraitsCallback&
+        worker_pool_index_for_traits_callback) {
+  std::unique_ptr<TaskSchedulerImpl> scheduler(
+      new TaskSchedulerImpl(worker_pool_index_for_traits_callback));
+  scheduler->Initialize(worker_pools);
   return scheduler;
 }
 
 TaskSchedulerImpl::~TaskSchedulerImpl() {
 #if DCHECK_IS_ON()
   DCHECK(join_for_testing_returned_.IsSignaled());
 #endif
 }
 
 void TaskSchedulerImpl::PostTaskWithTraits(
@@ skipping 15 matching lines @@
 
 void TaskSchedulerImpl::Shutdown() {
   // TODO(fdoray): Increase the priority of BACKGROUND tasks blocking shutdown.
   task_tracker_.Shutdown();
 }
 
 void TaskSchedulerImpl::JoinForTesting() {
 #if DCHECK_IS_ON()
   DCHECK(!join_for_testing_returned_.IsSignaled());
 #endif
-  background_worker_pool_->JoinForTesting();
-  background_file_io_worker_pool_->JoinForTesting();
-  normal_worker_pool_->JoinForTesting();
-  normal_file_io_worker_pool_->JoinForTesting();
+  for (const auto& worker_pool : worker_pools_)
+    worker_pool->JoinForTesting();
   service_thread_->JoinForTesting();
 #if DCHECK_IS_ON()
   join_for_testing_returned_.Signal();
 #endif
 }
 
-TaskSchedulerImpl::TaskSchedulerImpl()
+TaskSchedulerImpl::TaskSchedulerImpl(const WorkerPoolIndexForTraitsCallback&
+                                         worker_pool_index_for_traits_callback)
     : delayed_task_manager_(
-          Bind(&TaskSchedulerImpl::OnDelayedRunTimeUpdated, Unretained(this)))
+          Bind(&TaskSchedulerImpl::OnDelayedRunTimeUpdated, Unretained(this))),
+      worker_pool_index_for_traits_callback_(
+          worker_pool_index_for_traits_callback)
 #if DCHECK_IS_ON()
       ,
       join_for_testing_returned_(WaitableEvent::ResetPolicy::MANUAL,
                                  WaitableEvent::InitialState::NOT_SIGNALED)
 #endif
 {
+  DCHECK(!worker_pool_index_for_traits_callback_.is_null());
 }
 
-void TaskSchedulerImpl::Initialize() {
+void TaskSchedulerImpl::Initialize(
+    const std::vector<WorkerPoolCreationArgs>& worker_pools) {
+  DCHECK(!worker_pools.empty());
+
   using IORestriction = SchedulerWorkerPoolImpl::IORestriction;
 
   const SchedulerWorkerPoolImpl::ReEnqueueSequenceCallback
       re_enqueue_sequence_callback =
           Bind(&TaskSchedulerImpl::ReEnqueueSequenceCallback, Unretained(this));
 
-  // TODO(fdoray): Derive the number of threads per pool from hardware
-  // characteristics rather than using hard-coded constants.
-
-  // Passing pointers to objects owned by |this| to
-  // SchedulerWorkerPoolImpl::Create() is safe because a TaskSchedulerImpl can't
-  // be deleted before all its worker pools have been joined.
-  background_worker_pool_ = SchedulerWorkerPoolImpl::Create(
-      "TaskSchedulerBackground", ThreadPriority::BACKGROUND, 1U,
-      IORestriction::DISALLOWED, re_enqueue_sequence_callback, &task_tracker_,
-      &delayed_task_manager_);
-  CHECK(background_worker_pool_);
-
-  background_file_io_worker_pool_ = SchedulerWorkerPoolImpl::Create(
-      "TaskSchedulerBackgroundFileIO", ThreadPriority::BACKGROUND, 1U,
-      IORestriction::ALLOWED, re_enqueue_sequence_callback, &task_tracker_,
-      &delayed_task_manager_);
-  CHECK(background_file_io_worker_pool_);
-
-  normal_worker_pool_ = SchedulerWorkerPoolImpl::Create(
-      "TaskSchedulerForeground", ThreadPriority::NORMAL, 4U,
-      IORestriction::DISALLOWED, re_enqueue_sequence_callback, &task_tracker_,
-      &delayed_task_manager_);
-  CHECK(normal_worker_pool_);
-
-  normal_file_io_worker_pool_ = SchedulerWorkerPoolImpl::Create(
-      "TaskSchedulerForegroundFileIO", ThreadPriority::NORMAL, 12U,
-      IORestriction::ALLOWED, re_enqueue_sequence_callback, &task_tracker_,
-      &delayed_task_manager_);
-  CHECK(normal_file_io_worker_pool_);
+  for (const auto& worker_pool : worker_pools) {
+    // Passing pointers to objects owned by |this| to
+    // SchedulerWorkerPoolImpl::Create() is safe because a TaskSchedulerImpl
+    // can't be deleted before all its worker pools have been joined.
+    worker_pools_.push_back(SchedulerWorkerPoolImpl::Create(
+        worker_pool.name, worker_pool.thread_priority, worker_pool.max_threads,
+        worker_pool.io_restriction, re_enqueue_sequence_callback,
+        &task_tracker_, &delayed_task_manager_));
+    CHECK(worker_pools_.back());
+  }
 
   service_thread_ = SchedulerServiceThread::Create(&task_tracker_,
                                                    &delayed_task_manager_);
   CHECK(service_thread_);
 }
 
 SchedulerWorkerPool* TaskSchedulerImpl::GetWorkerPoolForTraits(
     const TaskTraits& traits) {
-  if (traits.with_file_io()) {
-    if (traits.priority() == TaskPriority::BACKGROUND)
-      return background_file_io_worker_pool_.get();
-    return normal_file_io_worker_pool_.get();
-  }
-
-  if (traits.priority() == TaskPriority::BACKGROUND)
-    return background_worker_pool_.get();
-  return normal_worker_pool_.get();
+  const size_t index = worker_pool_index_for_traits_callback_.Run(traits);
+  DCHECK_LT(index, worker_pools_.size());
+  return worker_pools_[index].get();
 }
 
 void TaskSchedulerImpl::ReEnqueueSequenceCallback(
     scoped_refptr<Sequence> sequence) {
   DCHECK(sequence);
 
   const SequenceSortKey sort_key = sequence->GetSortKey();
   TaskTraits traits(sequence->PeekTask()->traits);
 
   // Update the priority of |traits| so that the next task in |sequence| runs
   // with the highest priority in |sequence| as opposed to the next task's
   // specific priority.
   traits.WithPriority(sort_key.priority());
 
   GetWorkerPoolForTraits(traits)->ReEnqueueSequence(std::move(sequence),
                                                     sort_key);
 }
 
 void TaskSchedulerImpl::OnDelayedRunTimeUpdated() {
   service_thread_->WakeUp();
 }
 
 }  // namespace internal
 }  // namespace base