TaskScheduler: Make the worker pools of TaskSchedulerImpl configurable

base/task_scheduler/task_scheduler_impl.h

Index: base/task_scheduler/task_scheduler_impl.h
diff --git a/base/task_scheduler/task_scheduler_impl.h b/base/task_scheduler/task_scheduler_impl.h
index 4180d9faf5ef56299c638d13259e1dce20ba7d7c..02cbe3c6c51b29a0b08f89a9aa3f450ac0a128dc 100644
--- a/base/task_scheduler/task_scheduler_impl.h
+++ b/base/task_scheduler/task_scheduler_impl.h
@@ -5,32 +5,63 @@
 #ifndef BASE_TASK_SCHEDULER_TASK_SCHEDULER_IMPL_H_
 #define BASE_TASK_SCHEDULER_TASK_SCHEDULER_IMPL_H_
 
+#include <stddef.h>
+
 #include <memory>
+#include <string>
+#include <vector>
 
 #include "base/base_export.h"
+#include "base/callback.h"
 #include "base/logging.h"
 #include "base/macros.h"
 #include "base/memory/ref_counted.h"
 #include "base/synchronization/waitable_event.h"
 #include "base/task_runner.h"
 #include "base/task_scheduler/delayed_task_manager.h"
+#include "base/task_scheduler/scheduler_worker_pool_impl.h"
 #include "base/task_scheduler/sequence.h"
 #include "base/task_scheduler/task_scheduler.h"
 #include "base/task_scheduler/task_tracker.h"
 #include "base/task_scheduler/task_traits.h"
+#include "base/threading/thread.h"
 
 namespace base {
 namespace internal {
 
 class SchedulerServiceThread;
-class SchedulerWorkerPoolImpl;
 
 // Default TaskScheduler implementation. This class is thread-safe.
 class BASE_EXPORT TaskSchedulerImpl : public TaskScheduler {
  public:
-  // Creates and returns an initialized TaskSchedulerImpl. CHECKs on failure to
-  // do so (never returns null).
-  static std::unique_ptr<TaskSchedulerImpl> Create();
+  struct WorkerPoolCreationArgs {
+    // Name of the pool. Used to label the pool's threads.
+    std::string name;
+
+    // Priority of the pool's threads.
+    ThreadPriority thread_priority;
+
+    // Whether I/O is allowed in the pool.
+    SchedulerWorkerPoolImpl::IORestriction io_restriction;
+
+    // Maximum number of threads in the pool.
+    size_t max_threads;
+  };
+
+  // Returns the index of the worker pool in which a task with |traits| should
+  // run. This should be coded in a future-proof way: new traits should
+  // gracefully map to a default pool.
+  using WorkerPoolIndexForTraitsCallback =
+      Callback<size_t(const TaskTraits& traits)>;
+
+  // Creates and returns an initialized TaskSchedulerImpl. CHECKs on failure.
+  // |worker_pools| describes the worker pools to create.
+  // |worker_pool_index_for_traits_callback| returns the index in |worker_pools|
+  // of the worker pool in which a task with given traits should run.
+  static std::unique_ptr<TaskSchedulerImpl> Create(
+      const std::vector<WorkerPoolCreationArgs>& worker_pools,
+      const WorkerPoolIndexForTraitsCallback&
+          worker_pool_index_for_traits_callback);
 
   // Destroying a TaskSchedulerImpl is not allowed in production; it is always
   // leaked. In tests, it can only be destroyed after JoinForTesting() has
@@ -51,9 +82,10 @@ class BASE_EXPORT TaskSchedulerImpl : public TaskScheduler {
   void JoinForTesting();
 
  private:
-  TaskSchedulerImpl();
+  TaskSchedulerImpl(const WorkerPoolIndexForTraitsCallback&
+                        worker_pool_index_for_traits_callback);
 
-  void Initialize();
+  void Initialize(const std::vector<WorkerPoolCreationArgs>& worker_pools);
 
   // Returns the worker pool that runs Tasks with |traits|.
   SchedulerWorkerPool* GetWorkerPoolForTraits(const TaskTraits& traits);
@@ -68,19 +100,8 @@ class BASE_EXPORT TaskSchedulerImpl : public TaskScheduler {
 
   TaskTracker task_tracker_;
   DelayedTaskManager delayed_task_manager_;
-
-  // Worker pool for BACKGROUND Tasks without file I/O.
-  std::unique_ptr<SchedulerWorkerPoolImpl> background_worker_pool_;
-
-  // Worker pool for BACKGROUND Tasks with file I/O.
-  std::unique_ptr<SchedulerWorkerPoolImpl> background_file_io_worker_pool_;
-
-  // Worker pool for USER_VISIBLE and USER_BLOCKING Tasks without file I/O.
-  std::unique_ptr<SchedulerWorkerPoolImpl> normal_worker_pool_;
-
-  // Worker pool for USER_VISIBLE and USER_BLOCKING Tasks with file I/O.
-  std::unique_ptr<SchedulerWorkerPoolImpl> normal_file_io_worker_pool_;
-
+  const WorkerPoolIndexForTraitsCallback worker_pool_index_for_traits_callback_;
+  std::vector<std::unique_ptr<SchedulerWorkerPoolImpl>> worker_pools_;
   std::unique_ptr<SchedulerServiceThread> service_thread_;
 
 #if DCHECK_IS_ON()