base: Add TaskExecutors to support extension traits in post_task.h

Embedders can register a TaskExecutor to handle execution of Tasks
posted via the post_task.h API with specific extension TaskTraits.

The TaskScheduler becomes the default TaskExecutor.

Related patch introducing TaskTraits extensions:
https://chromium-review.googlesource.com/c/chromium/src/+/1169179
Design doc:
https://docs.google.com/document/d/1SGy9VTXUwyXEX_yBZ0ukFAnS8B0hDeMUJD-1iALaE-Q/edit?usp=sharing

TBR=pasko@chromium.org,primiano@chromium.org,imcheng@chromium.org,dominickn@chromium.org

Bug: 867421
Change-Id: I4083848431cca17aab3b6a411170ea32e5503776
Reviewed-on: https://chromium-review.googlesource.com/1175829Reviewed-by: Dominick Ng <dominickn@chromium.org>
Reviewed-by: François Doray <fdoray@chromium.org>
Reviewed-by: Gabriel Charette <gab@chromium.org>
Reviewed-by: Sami Kyöstilä <skyostil@chromium.org>
Commit-Queue: Eric Seckler <eseckler@chromium.org>
Cr-Commit-Position: refs/heads/master@{#584709}

base/BUILD.gn

@@ -871,6 +871,8 @@ jumbo_component("base") {
     "task/sequence_manager/work_queue_sets.cc",
     "task/sequence_manager/work_queue_sets.h",
     "task/single_thread_task_runner_thread_mode.h",
+    "task/task_executor.cc",
+    "task/task_executor.h",
     "task/task_scheduler/can_schedule_sequence_observer.h",
     "task/task_scheduler/delayed_task_manager.cc",
     "task/task_scheduler/delayed_task_manager.h",
@@ -2430,6 +2432,7 @@ test("base_unittests") {
     "system_monitor/system_monitor_unittest.cc",
     "task/cancelable_task_tracker_unittest.cc",
     "task/lazy_task_runner_unittest.cc",
+    "task/post_task_unittest.cc",
     "task/scoped_set_task_priority_for_current_thread_unittest.cc",
     "task/sequence_manager/intrusive_heap_unittest.cc",
     "task/sequence_manager/lazily_deallocated_deque_unittest.cc",

base/task/post_task.cc

@@ -8,6 +8,7 @@
 
 #include "base/logging.h"
 #include "base/task/scoped_set_task_priority_for_current_thread.h"
+#include "base/task/task_executor.h"
 #include "base/task/task_scheduler/task_scheduler.h"
 #include "base/threading/post_task_and_reply_impl.h"
 
@@ -40,92 +41,88 @@ TaskTraits GetTaskTraitsWithExplicitPriority(const TaskTraits& traits) {
                               {internal::GetTaskPriorityForCurrentThread()});
 }
 
+TaskExecutor* GetTaskExecutorForTraits(const TaskTraits& traits) {
+  DCHECK(TaskScheduler::GetInstance())
+      << "Ref. Prerequisite section of post_task.h.\n\n"
+         "Hint: if this is in a unit test, you're likely merely missing a "
+         "base::test::ScopedTaskEnvironment member in your fixture.\n";
+  TaskExecutor* executor = GetRegisteredTaskExecutorForTraits(traits);
+  return executor ? executor : TaskScheduler::GetInstance();
+}
+
 }  // namespace
 
-void PostTask(const Location& from_here, OnceClosure task) {
-  PostDelayedTask(from_here, std::move(task), TimeDelta());
+bool PostTask(const Location& from_here, OnceClosure task) {
+  return PostDelayedTask(from_here, std::move(task), TimeDelta());
 }
 
-void PostDelayedTask(const Location& from_here,
+bool PostDelayedTask(const Location& from_here,
                      OnceClosure task,
                      TimeDelta delay) {
-  PostDelayedTaskWithTraits(from_here, TaskTraits(), std::move(task), delay);
+  return PostDelayedTaskWithTraits(from_here, TaskTraits(), std::move(task),
+                                   delay);
 }
 
-void PostTaskAndReply(const Location& from_here,
+bool PostTaskAndReply(const Location& from_here,
                       OnceClosure task,
                       OnceClosure reply) {
-  PostTaskWithTraitsAndReply(from_here, TaskTraits(), std::move(task),
-                             std::move(reply));
+  return PostTaskWithTraitsAndReply(from_here, TaskTraits(), std::move(task),
+                                    std::move(reply));
 }
 
-void PostTaskWithTraits(const Location& from_here,
+bool PostTaskWithTraits(const Location& from_here,
                         const TaskTraits& traits,
                         OnceClosure task) {
-  PostDelayedTaskWithTraits(from_here, traits, std::move(task), TimeDelta());
+  return PostDelayedTaskWithTraits(from_here, traits, std::move(task),
+                                   TimeDelta());
 }
 
-void PostDelayedTaskWithTraits(const Location& from_here,
+bool PostDelayedTaskWithTraits(const Location& from_here,
                                const TaskTraits& traits,
                                OnceClosure task,
                                TimeDelta delay) {
-  DCHECK(TaskScheduler::GetInstance())
-      << "Ref. Prerequisite section of post_task.h.\n\n"
-         "Hint: if this is in a unit test, you're likely merely missing a "
-         "base::test::ScopedTaskEnvironment member in your fixture.\n";
-  TaskScheduler::GetInstance()->PostDelayedTaskWithTraits(
-      from_here, GetTaskTraitsWithExplicitPriority(traits), std::move(task),
-      std::move(delay));
+  const TaskTraits adjusted_traits = GetTaskTraitsWithExplicitPriority(traits);
+  return GetTaskExecutorForTraits(adjusted_traits)
+      ->PostDelayedTaskWithTraits(from_here, adjusted_traits, std::move(task),
+                                  delay);
 }
 
-void PostTaskWithTraitsAndReply(const Location& from_here,
+bool PostTaskWithTraitsAndReply(const Location& from_here,
                                 const TaskTraits& traits,
                                 OnceClosure task,
                                 OnceClosure reply) {
-  PostTaskAndReplyWithTraitsTaskRunner(traits).PostTaskAndReply(
+  return PostTaskAndReplyWithTraitsTaskRunner(traits).PostTaskAndReply(
       from_here, std::move(task), std::move(reply));
 }
 
 scoped_refptr<TaskRunner> CreateTaskRunnerWithTraits(const TaskTraits& traits) {
-  DCHECK(TaskScheduler::GetInstance())
-      << "Ref. Prerequisite section of post_task.h.\n\n"
-         "Hint: if this is in a unit test, you're likely merely missing a "
-         "base::test::ScopedTaskEnvironment member in your fixture.\n";
-  return TaskScheduler::GetInstance()->CreateTaskRunnerWithTraits(
-      GetTaskTraitsWithExplicitPriority(traits));
+  const TaskTraits adjusted_traits = GetTaskTraitsWithExplicitPriority(traits);
+  return GetTaskExecutorForTraits(adjusted_traits)
+      ->CreateTaskRunnerWithTraits(adjusted_traits);
 }
 
 scoped_refptr<SequencedTaskRunner> CreateSequencedTaskRunnerWithTraits(
     const TaskTraits& traits) {
-  DCHECK(TaskScheduler::GetInstance())
-      << "Ref. Prerequisite section of post_task.h.\n\n"
-         "Hint: if this is in a unit test, you're likely merely missing a "
-         "base::test::ScopedTaskEnvironment member in your fixture.\n";
-  return TaskScheduler::GetInstance()->CreateSequencedTaskRunnerWithTraits(
-      GetTaskTraitsWithExplicitPriority(traits));
+  const TaskTraits adjusted_traits = GetTaskTraitsWithExplicitPriority(traits);
+  return GetTaskExecutorForTraits(adjusted_traits)
+      ->CreateSequencedTaskRunnerWithTraits(adjusted_traits);
 }
 
 scoped_refptr<SingleThreadTaskRunner> CreateSingleThreadTaskRunnerWithTraits(
     const TaskTraits& traits,
     SingleThreadTaskRunnerThreadMode thread_mode) {
-  DCHECK(TaskScheduler::GetInstance())
-      << "Ref. Prerequisite section of post_task.h.\n\n"
-         "Hint: if this is in a unit test, you're likely merely missing a "
-         "base::test::ScopedTaskEnvironment member in your fixture.\n";
-  return TaskScheduler::GetInstance()->CreateSingleThreadTaskRunnerWithTraits(
-      GetTaskTraitsWithExplicitPriority(traits), thread_mode);
+  const TaskTraits adjusted_traits = GetTaskTraitsWithExplicitPriority(traits);
+  return GetTaskExecutorForTraits(adjusted_traits)
+      ->CreateSingleThreadTaskRunnerWithTraits(adjusted_traits, thread_mode);
 }
 
 #if defined(OS_WIN)
 scoped_refptr<SingleThreadTaskRunner> CreateCOMSTATaskRunnerWithTraits(
     const TaskTraits& traits,
     SingleThreadTaskRunnerThreadMode thread_mode) {
-  DCHECK(TaskScheduler::GetInstance())
-      << "Ref. Prerequisite section of post_task.h.\n\n"
-         "Hint: if this is in a unit test, you're likely merely missing a "
-         "base::test::ScopedTaskEnvironment member in your fixture.\n";
-  return TaskScheduler::GetInstance()->CreateCOMSTATaskRunnerWithTraits(
-      GetTaskTraitsWithExplicitPriority(traits), thread_mode);
+  const TaskTraits adjusted_traits = GetTaskTraitsWithExplicitPriority(traits);
+  return GetTaskExecutorForTraits(adjusted_traits)
+      ->CreateCOMSTATaskRunnerWithTraits(adjusted_traits, thread_mode);
 }
 #endif  // defined(OS_WIN)
 

base/task/post_task_unittest.cc

+// Copyright 2018 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/post_task.h"
+
+#include "base/bind_helpers.h"
+#include "base/task/task_executor.h"
+#include "base/task/test_task_traits_extension.h"
+#include "base/test/scoped_task_environment.h"
+#include "base/test/test_simple_task_runner.h"
+#include "build/build_config.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+using ::testing::_;
+using ::testing::Invoke;
+using ::testing::Return;
+
+namespace base {
+
+namespace {
+
+class MockTaskExecutor : public TaskExecutor {
+ public:
+  MockTaskExecutor() {
+    ON_CALL(*this, PostDelayedTaskWithTraitsMock(_, _, _, _))
+        .WillByDefault(Invoke([this](const Location& from_here,
+                                     const TaskTraits& traits,
+                                     OnceClosure& task, TimeDelta delay) {
+          return runner_->PostDelayedTask(from_here, std::move(task), delay);
+        }));
+    ON_CALL(*this, CreateTaskRunnerWithTraits(_))
+        .WillByDefault(Return(runner_));
+    ON_CALL(*this, CreateSequencedTaskRunnerWithTraits(_))
+        .WillByDefault(Return(runner_));
+    ON_CALL(*this, CreateSingleThreadTaskRunnerWithTraits(_, _))
+        .WillByDefault(Return(runner_));
+#if defined(OS_WIN)
+    ON_CALL(*this, CreateCOMSTATaskRunnerWithTraits(_, _))
+        .WillByDefault(Return(runner_));
+#endif  // defined(OS_WIN)
+  }
+
+  // TaskExecutor:
+  // Helper because gmock doesn't support move-only types.
+  bool PostDelayedTaskWithTraits(const Location& from_here,
+                                 const TaskTraits& traits,
+                                 OnceClosure task,
+                                 TimeDelta delay) override {
+    return PostDelayedTaskWithTraitsMock(from_here, traits, task, delay);
+  }
+  MOCK_METHOD4(PostDelayedTaskWithTraitsMock,
+               bool(const Location& from_here,
+                    const TaskTraits& traits,
+                    OnceClosure& task,
+                    TimeDelta delay));
+  MOCK_METHOD1(CreateTaskRunnerWithTraits,
+               scoped_refptr<TaskRunner>(const TaskTraits& traits));
+  MOCK_METHOD1(CreateSequencedTaskRunnerWithTraits,
+               scoped_refptr<SequencedTaskRunner>(const TaskTraits& traits));
+  MOCK_METHOD2(CreateSingleThreadTaskRunnerWithTraits,
+               scoped_refptr<SingleThreadTaskRunner>(
+                   const TaskTraits& traits,
+                   SingleThreadTaskRunnerThreadMode thread_mode));
+#if defined(OS_WIN)
+  MOCK_METHOD2(CreateCOMSTATaskRunnerWithTraits,
+               scoped_refptr<SingleThreadTaskRunner>(
+                   const TaskTraits& traits,
+                   SingleThreadTaskRunnerThreadMode thread_mode));
+#endif  // defined(OS_WIN)
+
+  TestSimpleTaskRunner* runner() const { return runner_.get(); }
+
+ private:
+  scoped_refptr<TestSimpleTaskRunner> runner_ =
+      MakeRefCounted<TestSimpleTaskRunner>();
+
+  DISALLOW_COPY_AND_ASSIGN(MockTaskExecutor);
+};
+
+}  // namespace
+
+class PostTaskTestWithExecutor : public ::testing::Test {
+ public:
+  void SetUp() override {
+    RegisterTaskExecutor(TestTaskTraitsExtension::kExtensionId, &executor_);
+  }
+
+  void TearDown() override {
+    UnregisterTaskExecutorForTesting(TestTaskTraitsExtension::kExtensionId);
+  }
+
+ protected:
+  testing::StrictMock<MockTaskExecutor> executor_;
+  test::ScopedTaskEnvironment scoped_task_environment_;
+};
+
+TEST_F(PostTaskTestWithExecutor, PostTaskToTaskScheduler) {
+  // Tasks without extension should not go to the TestTaskExecutor.
+  EXPECT_TRUE(PostTask(FROM_HERE, DoNothing()));
+  EXPECT_FALSE(executor_.runner()->HasPendingTask());
+
+  EXPECT_TRUE(PostTaskWithTraits(FROM_HERE, {MayBlock()}, DoNothing()));
+  EXPECT_FALSE(executor_.runner()->HasPendingTask());
+
+  // Task runners without extension should not be the executor's.
+  auto task_runner = CreateTaskRunnerWithTraits({});
+  EXPECT_NE(executor_.runner(), task_runner);
+  auto sequenced_task_runner = CreateSequencedTaskRunnerWithTraits({});
+  EXPECT_NE(executor_.runner(), sequenced_task_runner);
+  auto single_thread_task_runner = CreateSingleThreadTaskRunnerWithTraits({});
+  EXPECT_NE(executor_.runner(), single_thread_task_runner);
+#if defined(OS_WIN)
+  auto comsta_task_runner = CreateCOMSTATaskRunnerWithTraits({});
+  EXPECT_NE(executor_.runner(), comsta_task_runner);
+#endif  // defined(OS_WIN)
+}
+
+TEST_F(PostTaskTestWithExecutor, PostTaskToTaskExecutor) {
+  // Tasks with extension should go to the executor.
+  {
+    TaskTraits traits = {TestExtensionBoolTrait()};
+    TaskTraits traits_with_explicit_priority =
+        TaskTraits::Override(traits, {TaskPriority::USER_VISIBLE});
+    EXPECT_CALL(executor_, PostDelayedTaskWithTraitsMock(
+                               _, traits_with_explicit_priority, _, _))
+        .Times(1);
+    EXPECT_TRUE(PostTaskWithTraits(FROM_HERE, traits, DoNothing()));
+    EXPECT_TRUE(executor_.runner()->HasPendingTask());
+    executor_.runner()->ClearPendingTasks();
+  }
+
+  {
+    TaskTraits traits = {MayBlock(), TestExtensionBoolTrait()};
+    TaskTraits traits_with_explicit_priority =
+        TaskTraits::Override(traits, {TaskPriority::USER_VISIBLE});
+    EXPECT_CALL(executor_, PostDelayedTaskWithTraitsMock(
+                               _, traits_with_explicit_priority, _, _))
+        .Times(1);
+    EXPECT_TRUE(PostTaskWithTraits(FROM_HERE, traits, DoNothing()));
+    EXPECT_TRUE(executor_.runner()->HasPendingTask());
+    executor_.runner()->ClearPendingTasks();
+  }
+
+  {
+    TaskTraits traits = {TestExtensionEnumTrait::kB, TestExtensionBoolTrait()};
+    TaskTraits traits_with_explicit_priority =
+        TaskTraits::Override(traits, {TaskPriority::USER_VISIBLE});
+    EXPECT_CALL(executor_, PostDelayedTaskWithTraitsMock(
+                               _, traits_with_explicit_priority, _, _))
+        .Times(1);
+    EXPECT_TRUE(PostTaskWithTraits(FROM_HERE, traits, DoNothing()));
+    EXPECT_TRUE(executor_.runner()->HasPendingTask());
+    executor_.runner()->ClearPendingTasks();
+  }
+
+  // Task runners with extension should be the executor's.
+  {
+    TaskTraits traits = {TestExtensionBoolTrait()};
+    TaskTraits traits_with_explicit_priority =
+        TaskTraits::Override(traits, {TaskPriority::USER_VISIBLE});
+    EXPECT_CALL(executor_,
+                CreateTaskRunnerWithTraits(traits_with_explicit_priority))
+        .Times(1);
+    auto task_runner = CreateTaskRunnerWithTraits(traits);
+    EXPECT_EQ(executor_.runner(), task_runner);
+    EXPECT_CALL(executor_, CreateSequencedTaskRunnerWithTraits(
+                               traits_with_explicit_priority))
+        .Times(1);
+    auto sequenced_task_runner = CreateSequencedTaskRunnerWithTraits(traits);
+    EXPECT_EQ(executor_.runner(), sequenced_task_runner);
+    EXPECT_CALL(executor_, CreateSingleThreadTaskRunnerWithTraits(
+                               traits_with_explicit_priority, _))
+        .Times(1);
+    auto single_thread_task_runner =
+        CreateSingleThreadTaskRunnerWithTraits(traits);
+    EXPECT_EQ(executor_.runner(), single_thread_task_runner);
+#if defined(OS_WIN)
+    EXPECT_CALL(executor_, CreateCOMSTATaskRunnerWithTraits(
+                               traits_with_explicit_priority, _))
+        .Times(1);
+    auto comsta_task_runner = CreateCOMSTATaskRunnerWithTraits(traits);
+    EXPECT_EQ(executor_.runner(), comsta_task_runner);
+#endif  // defined(OS_WIN)
+  }
+}
+
+}  // namespace base

base/task/task_executor.cc

+// Copyright 2018 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/task_executor.h"
+
+#include "base/no_destructor.h"
+#include "base/task/task_traits.h"
+#include "base/task/task_traits_extension.h"
+
+namespace base {
+
+namespace {
+
+// Maps TaskTraits extension IDs to registered TaskExecutors. Index |n|
+// corresponds to id |n - 1|.
+using TaskExecutorMap =
+    std::array<TaskExecutor*, TaskTraitsExtensionStorage::kMaxExtensionId>;
+TaskExecutorMap* GetTaskExecutorMap() {
+  static NoDestructor<TaskExecutorMap> executors(TaskExecutorMap{});
+  return executors.get();
+}
+
+static_assert(
+    TaskTraitsExtensionStorage::kInvalidExtensionId == 0,
+    "TaskExecutorMap depends on 0 being an invalid TaskTraits extension ID");
+
+}  // namespace
+
+void RegisterTaskExecutor(uint8_t extension_id, TaskExecutor* task_executor) {
+  DCHECK_NE(extension_id, TaskTraitsExtensionStorage::kInvalidExtensionId);
+  DCHECK_LE(extension_id, TaskTraitsExtensionStorage::kMaxExtensionId);
+  DCHECK_EQ((*GetTaskExecutorMap())[extension_id - 1], nullptr);
+
+  (*GetTaskExecutorMap())[extension_id - 1] = task_executor;
+}
+
+void UnregisterTaskExecutorForTesting(uint8_t extension_id) {
+  DCHECK_NE(extension_id, TaskTraitsExtensionStorage::kInvalidExtensionId);
+  DCHECK_LE(extension_id, TaskTraitsExtensionStorage::kMaxExtensionId);
+  DCHECK_NE((*GetTaskExecutorMap())[extension_id - 1], nullptr);
+
+  (*GetTaskExecutorMap())[extension_id - 1] = nullptr;
+}
+
+TaskExecutor* GetRegisteredTaskExecutorForTraits(const TaskTraits& traits) {
+  uint8_t extension_id = traits.extension_id();
+  if (extension_id != TaskTraitsExtensionStorage::kInvalidExtensionId) {
+    TaskExecutor* executor = (*GetTaskExecutorMap())[extension_id - 1];
+    DCHECK(executor);
+    return executor;
+  }
+
+  return nullptr;
+}
+
+}  // namespace base
\ No newline at end of file

base/task/task_executor.h

+// Copyright 2018 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.
+
+#ifndef BASE_TASK_TASK_EXECUTOR_H_
+#define BASE_TASK_TASK_EXECUTOR_H_
+
+#include <stdint.h>
+
+#include "base/base_export.h"
+#include "base/memory/ref_counted.h"
+#include "base/sequenced_task_runner.h"
+#include "base/single_thread_task_runner.h"
+#include "base/task/single_thread_task_runner_thread_mode.h"
+#include "base/task_runner.h"
+#include "build/build_config.h"
+
+namespace base {
+
+class Location;
+class TaskTraits;
+
+// A TaskExecutor can execute Tasks with a specific TaskTraits extension id. To
+// handle Tasks posted via the //base/task/post_task.h API, the TaskExecutor
+// should be registered by calling RegisterTaskExecutor().
+class BASE_EXPORT TaskExecutor {
+ public:
+  virtual ~TaskExecutor() = default;
+
+  // Posts |task| with a |delay| and specific |traits|. |delay| can be zero. For
+  // one off tasks that don't require a TaskRunner. Returns false if the task
+  // definitely won't run because of current shutdown state.
+  virtual bool PostDelayedTaskWithTraits(const Location& from_here,
+                                         const TaskTraits& traits,
+                                         OnceClosure task,
+                                         TimeDelta delay) = 0;
+
+  // Returns a TaskRunner whose PostTask invocations result in scheduling tasks
+  // using |traits|. Tasks may run in any order and in parallel.
+  virtual scoped_refptr<TaskRunner> CreateTaskRunnerWithTraits(
+      const TaskTraits& traits) = 0;
+
+  // Returns a SequencedTaskRunner whose PostTask invocations result in
+  // scheduling tasks using |traits|. Tasks run one at a time in posting order.
+  virtual scoped_refptr<SequencedTaskRunner>
+  CreateSequencedTaskRunnerWithTraits(const TaskTraits& traits) = 0;
+
+  // Returns a SingleThreadTaskRunner whose PostTask invocations result in
+  // scheduling tasks using |traits|. Tasks run on a single thread in posting
+  // order. If |traits| identifies an existing thread,
+  // SingleThreadTaskRunnerThreadMode::SHARED must be used.
+  virtual scoped_refptr<SingleThreadTaskRunner>
+  CreateSingleThreadTaskRunnerWithTraits(
+      const TaskTraits& traits,
+      SingleThreadTaskRunnerThreadMode thread_mode) = 0;
+
+#if defined(OS_WIN)
+  // Returns a SingleThreadTaskRunner whose PostTask invocations result in
+  // scheduling tasks using |traits| in a COM Single-Threaded Apartment. Tasks
+  // run in the same Single-Threaded Apartment in posting order for the returned
+  // SingleThreadTaskRunner. If |traits| identifies an existing thread,
+  // SingleThreadTaskRunnerThreadMode::SHARED must be used.
+  virtual scoped_refptr<SingleThreadTaskRunner>
+  CreateCOMSTATaskRunnerWithTraits(
+      const TaskTraits& traits,
+      SingleThreadTaskRunnerThreadMode thread_mode) = 0;
+#endif  // defined(OS_WIN)
+};
+
+// Register a TaskExecutor with the //base/task/post_task.h API in the current
+// process for tasks subsequently posted with a TaskTraits extension with the
+// given |extension_id|. All executors need to be registered before any tasks
+// are posted with |extension_id|. Only one executor per |extension_id| is
+// supported.
+void BASE_EXPORT RegisterTaskExecutor(uint8_t extension_id,
+                                      TaskExecutor* task_executor);
+void BASE_EXPORT UnregisterTaskExecutorForTesting(uint8_t extension_id);
+
+// Determines whether a registered TaskExecutor will handle tasks with the given
+// |traits| and, if so, returns a pointer to it. Otherwise, returns |nullptr|.
+TaskExecutor* GetRegisteredTaskExecutorForTraits(const TaskTraits& traits);
+
+}  // namespace base
+
+#endif  // BASE_TASK_TASK_EXECUTOR_H_

base/task/task_scheduler/task_scheduler.h

@@ -16,6 +16,7 @@
 #include "base/single_thread_task_runner.h"
 #include "base/strings/string_piece.h"
 #include "base/task/single_thread_task_runner_thread_mode.h"
+#include "base/task/task_executor.h"
 #include "base/task/task_scheduler/scheduler_worker_pool_params.h"
 #include "base/task/task_traits.h"
 #include "base/task_runner.h"
@@ -35,7 +36,6 @@ class BrowserMainLoopTest_CreateThreadsInSingleProcess_Test;
 namespace base {
 
 class HistogramBase;
-class Location;
 class SchedulerWorkerObserver;
 
 // Interface for a task scheduler and static methods to manage the instance used
@@ -49,7 +49,7 @@ class SchedulerWorkerObserver;
 // Note: All TaskScheduler users should go through base/task/post_task.h instead
 // of this interface except for the one callsite per process which manages the
 // process's instance.
-class BASE_EXPORT TaskScheduler {
+class BASE_EXPORT TaskScheduler : public TaskExecutor {
  public:
   struct BASE_EXPORT InitParams {
     enum class SharedWorkerPoolEnvironment {
@@ -82,7 +82,7 @@ class BASE_EXPORT TaskScheduler {
   // Destroying a TaskScheduler is not allowed in production; it is always
   // leaked. In tests, it should only be destroyed after JoinForTesting() has
   // returned.
-  virtual ~TaskScheduler() = default;
+  ~TaskScheduler() override = default;
 
   // Allows the task scheduler to create threads and run tasks following the
   // |init_params| specification.
@@ -96,46 +96,6 @@ class BASE_EXPORT TaskScheduler {
       const InitParams& init_params,
       SchedulerWorkerObserver* scheduler_worker_observer = nullptr) = 0;
 
-  // Posts |task| with a |delay| and specific |traits|. |delay| can be zero.
-  // For one off tasks that don't require a TaskRunner.
-  virtual void PostDelayedTaskWithTraits(const Location& from_here,
-                                         const TaskTraits& traits,
-                                         OnceClosure task,
-                                         TimeDelta delay) = 0;
-
-  // Returns a TaskRunner whose PostTask invocations result in scheduling tasks
-  // using |traits|. Tasks may run in any order and in parallel.
-  virtual scoped_refptr<TaskRunner> CreateTaskRunnerWithTraits(
-      const TaskTraits& traits) = 0;
-
-  // Returns a SequencedTaskRunner whose PostTask invocations result in
-  // scheduling tasks using |traits|. Tasks run one at a time in posting order.
-  virtual scoped_refptr<SequencedTaskRunner>
-  CreateSequencedTaskRunnerWithTraits(const TaskTraits& traits) = 0;
-
-  // Returns a SingleThreadTaskRunner whose PostTask invocations result in
-  // scheduling tasks using |traits|. Tasks run on a single thread in posting
-  // order.
-  virtual scoped_refptr<SingleThreadTaskRunner>
-  CreateSingleThreadTaskRunnerWithTraits(
-      const TaskTraits& traits,
-      SingleThreadTaskRunnerThreadMode thread_mode) = 0;
-
-#if defined(OS_WIN)
-  // Returns a SingleThreadTaskRunner whose PostTask invocations result in
-  // scheduling tasks using |traits| in a COM Single-Threaded Apartment. Tasks
-  // run in the same Single-Threaded Apartment in posting order for the returned
-  // SingleThreadTaskRunner. There is not necessarily a one-to-one
-  // correspondence between SingleThreadTaskRunners and Single-Threaded
-  // Apartments. The implementation is free to share apartments or create new
-  // apartments as necessary. In either case, care should be taken to make sure
-  // COM pointers are not smuggled across apartments.
-  virtual scoped_refptr<SingleThreadTaskRunner>
-  CreateCOMSTATaskRunnerWithTraits(
-      const TaskTraits& traits,
-      SingleThreadTaskRunnerThreadMode thread_mode) = 0;
-#endif  // defined(OS_WIN)
-
   // Returns a vector of all histograms available in this task scheduler.
   virtual std::vector<const HistogramBase*> GetHistograms() const = 0;
 

base/task/task_scheduler/task_scheduler_impl.cc

@@ -177,13 +177,13 @@ void TaskSchedulerImpl::Start(
   }
 }
 
-void TaskSchedulerImpl::PostDelayedTaskWithTraits(const Location& from_here,
+bool TaskSchedulerImpl::PostDelayedTaskWithTraits(const Location& from_here,
                                                   const TaskTraits& traits,
                                                   OnceClosure task,
                                                   TimeDelta delay) {
   // Post |task| as part of a one-off single-task Sequence.
   const TaskTraits new_traits = SetUserBlockingPriorityIfNeeded(traits);
-  GetWorkerPoolForTraits(new_traits)
+  return GetWorkerPoolForTraits(new_traits)
       ->PostTaskWithSequence(
           Task(from_here, std::move(task), new_traits, delay),
           MakeRefCounted<Sequence>());

base/task/task_scheduler/task_scheduler_impl.h

@@ -64,7 +64,16 @@ class BASE_EXPORT TaskSchedulerImpl : public TaskScheduler {
   // TaskScheduler:
   void Start(const TaskScheduler::InitParams& init_params,
              SchedulerWorkerObserver* scheduler_worker_observer) override;
-  void PostDelayedTaskWithTraits(const Location& from_here,
+  std::vector<const HistogramBase*> GetHistograms() const override;
+  int GetMaxConcurrentNonBlockedTasksWithTraitsDeprecated(
+      const TaskTraits& traits) const override;
+  void Shutdown() override;
+  void FlushForTesting() override;
+  void FlushAsyncForTesting(OnceClosure flush_callback) override;
+  void JoinForTesting() override;
+
+  // TaskExecutor:
+  bool PostDelayedTaskWithTraits(const Location& from_here,
                                  const TaskTraits& traits,
                                  OnceClosure task,
                                  TimeDelta delay) override;
@@ -80,13 +89,6 @@ class BASE_EXPORT TaskSchedulerImpl : public TaskScheduler {
       const TaskTraits& traits,
       SingleThreadTaskRunnerThreadMode thread_mode) override;
 #endif  // defined(OS_WIN)
-  std::vector<const HistogramBase*> GetHistograms() const override;
-  int GetMaxConcurrentNonBlockedTasksWithTraitsDeprecated(
-      const TaskTraits& traits) const override;
-  void Shutdown() override;
-  void FlushForTesting() override;
-  void FlushAsyncForTesting(OnceClosure flush_callback) override;
-  void JoinForTesting() override;
 
  private:
   // Returns the worker pool that runs Tasks with |traits|.

base/task/task_traits.h

@@ -156,20 +156,20 @@ class BASE_EXPORT TaskTraits {
             class CheckArgumentsAreValidBaseTraits = std::enable_if_t<
                 trait_helpers::AreValidTraits<ValidTrait, ArgTypes...>::value>>
   constexpr TaskTraits(ArgTypes... args)
-      : priority_set_explicitly_(
-            trait_helpers::HasArgOfType<TaskPriority, ArgTypes...>::value),
-        priority_(trait_helpers::GetValueFromArgList(
+      : priority_(trait_helpers::GetValueFromArgList(
             trait_helpers::EnumArgGetter<TaskPriority,
                                          TaskPriority::USER_VISIBLE>(),
             args...)),
-        shutdown_behavior_set_explicitly_(
-            trait_helpers::HasArgOfType<TaskShutdownBehavior,
-                                        ArgTypes...>::value),
         shutdown_behavior_(trait_helpers::GetValueFromArgList(
             trait_helpers::EnumArgGetter<
                 TaskShutdownBehavior,
                 TaskShutdownBehavior::SKIP_ON_SHUTDOWN>(),
             args...)),
+        priority_set_explicitly_(
+            trait_helpers::HasArgOfType<TaskPriority, ArgTypes...>::value),
+        shutdown_behavior_set_explicitly_(
+            trait_helpers::HasArgOfType<TaskShutdownBehavior,
+                                        ArgTypes...>::value),
         may_block_(trait_helpers::GetValueFromArgList(
             trait_helpers::BooleanArgGetter<MayBlock>(),
             args...)),
@@ -200,6 +200,20 @@ class BASE_EXPORT TaskTraits {
   constexpr TaskTraits(const TaskTraits& other) = default;
   TaskTraits& operator=(const TaskTraits& other) = default;
 
+  // TODO(eseckler): Default the comparison operator once C++20 arrives.
+  bool operator==(const TaskTraits& other) const {
+    static_assert(24 == sizeof(TaskTraits),
+                  "Update comparison operator when TaskTraits change");
+    return priority_set_explicitly_ == other.priority_set_explicitly_ &&
+           priority_ == other.priority_ &&
+           shutdown_behavior_set_explicitly_ ==
+               other.shutdown_behavior_set_explicitly_ &&
+           shutdown_behavior_ == other.shutdown_behavior_ &&
+           may_block_ == other.may_block_ &&
+           with_base_sync_primitives_ == other.with_base_sync_primitives_ &&
+           extension_ == other.extension_;
+  }
+
   // Returns TaskTraits constructed by combining |left| and |right|. If a trait
   // is specified in both |left| and |right|, the returned TaskTraits will have
   // the value from |right|. Note that extension traits are not merged: any
@@ -248,23 +262,23 @@ class BASE_EXPORT TaskTraits {
 
  private:
   constexpr TaskTraits(const TaskTraits& left, const TaskTraits& right)
-      : priority_set_explicitly_(left.priority_set_explicitly_ ||
-                                 right.priority_set_explicitly_),
+      : extension_(right.extension_.extension_id !=
+                           TaskTraitsExtensionStorage::kInvalidExtensionId
+                       ? right.extension_
+                       : left.extension_),
         priority_(right.priority_set_explicitly_ ? right.priority_
                                                  : left.priority_),
-        shutdown_behavior_set_explicitly_(
-            left.shutdown_behavior_set_explicitly_ ||
-            right.shutdown_behavior_set_explicitly_),
         shutdown_behavior_(right.shutdown_behavior_set_explicitly_
                                ? right.shutdown_behavior_
                                : left.shutdown_behavior_),
+        priority_set_explicitly_(left.priority_set_explicitly_ ||
+                                 right.priority_set_explicitly_),
+        shutdown_behavior_set_explicitly_(
+            left.shutdown_behavior_set_explicitly_ ||
+            right.shutdown_behavior_set_explicitly_),
         may_block_(left.may_block_ || right.may_block_),
         with_base_sync_primitives_(left.with_base_sync_primitives_ ||
-                                   right.with_base_sync_primitives_),
-        extension_(right.extension_.extension_id !=
-                           TaskTraitsExtensionStorage::kInvalidExtensionId
-                       ? right.extension_
-                       : left.extension_) {}
+                                   right.with_base_sync_primitives_) {}
 
   // Helper constructor which selects those arguments from |args| that are
   // indicated by the index_sequence and forwards them to the public
@@ -276,14 +290,14 @@ class BASE_EXPORT TaskTraits {
             std::get<Indices>(std::forward<std::tuple<ArgTypes...>>(args))...) {
   }
 
-  bool priority_set_explicitly_;
+  // Ordered for packing.
+  TaskTraitsExtensionStorage extension_;
   TaskPriority priority_;
-  bool shutdown_behavior_set_explicitly_;
   TaskShutdownBehavior shutdown_behavior_;
+  bool priority_set_explicitly_;
+  bool shutdown_behavior_set_explicitly_;
   bool may_block_;
   bool with_base_sync_primitives_;
-
-  TaskTraitsExtensionStorage extension_;
 };
 
 // Returns string literals for the enums defined in this file. These methods

base/task/task_traits_extension.h

@@ -134,11 +134,15 @@ struct BASE_EXPORT TaskTraitsExtensionStorage {
   inline TaskTraitsExtensionStorage& operator=(
       const TaskTraitsExtensionStorage& other) = default;
 
+  inline bool operator==(const TaskTraitsExtensionStorage& other) const;
+
   enum ExtensionId : uint8_t {
     kInvalidExtensionId = 0,
     // The embedder is responsible for assigning the remaining values uniquely.
     kFirstEmbedderExtensionId = 1,
-    kMaxExtensionId = 255
+    // Maximum number of extension types is artificially limited to support
+    // super efficient TaskExecutor lookup in post_task.cc.
+    kMaxExtensionId = 4
   };
 
   // Identifies the type of extension. See ExtensionId enum above.
@@ -167,6 +171,15 @@ inline constexpr TaskTraitsExtensionStorage::TaskTraitsExtensionStorage(
 inline constexpr TaskTraitsExtensionStorage::TaskTraitsExtensionStorage(
     const TaskTraitsExtensionStorage& other) = default;
 
+// TODO(eseckler): Default the comparison operator once C++20 arrives.
+inline bool TaskTraitsExtensionStorage::operator==(
+    const TaskTraitsExtensionStorage& other) const {
+  static_assert(
+      9 == sizeof(TaskTraitsExtensionStorage),
+      "Update comparison operator when TaskTraitsExtensionStorage changes");
+  return extension_id == other.extension_id && data == other.data;
+}
+
 // Default implementation of MakeTaskTraitsExtension template function, which
 // doesn't accept any traits and does not create an extension.
 template <class Unused = void>

chrome/browser/ui/media_router/media_router_file_dialog.cc

@@ -130,9 +130,8 @@ MediaRouterFileDialog::MediaRouterFileDialog(
 MediaRouterFileDialog::MediaRouterFileDialog(
     MediaRouterFileDialogDelegate* delegate,
     std::unique_ptr<FileSystemDelegate> file_system_delegate)
-    : task_runner_(
-          base::TaskScheduler::GetInstance()->CreateTaskRunnerWithTraits(
-              {base::MayBlock(), base::TaskPriority::USER_VISIBLE})),
+    : task_runner_(base::CreateTaskRunnerWithTraits(
+          {base::MayBlock(), base::TaskPriority::USER_VISIBLE})),
       file_system_delegate_(std::move(file_system_delegate)),
       delegate_(delegate) {}
 

chrome/browser/web_applications/components/web_app_shortcut.cc

@@ -7,7 +7,6 @@
 #include "base/i18n/file_util_icu.h"
 #include "base/strings/utf_string_conversions.h"
 #include "base/task/post_task.h"
-#include "base/task/task_scheduler/task_scheduler.h"
 #include "build/build_config.h"
 #include "chrome/browser/web_applications/components/web_app_helpers.h"
 #include "chrome/common/chrome_constants.h"
@@ -105,10 +104,10 @@ scoped_refptr<base::TaskRunner> GetShortcutIOTaskRunner() {
       base::TaskShutdownBehavior::BLOCK_SHUTDOWN};
 
 #if defined(OS_WIN)
-  return base::TaskScheduler::GetInstance()->CreateCOMSTATaskRunnerWithTraits(
+  return base::CreateCOMSTATaskRunnerWithTraits(
       traits, base::SingleThreadTaskRunnerThreadMode::SHARED);
 #else
-  return base::TaskScheduler::GetInstance()->CreateTaskRunnerWithTraits(traits);
+  return base::CreateTaskRunnerWithTraits(traits);
 #endif
 }
 

content/browser/tracing/cast_tracing_agent.cc

@@ -9,7 +9,7 @@
 #include "base/bind.h"
 #include "base/bind_helpers.h"
 #include "base/logging.h"
-#include "base/task/task_scheduler/task_scheduler.h"
+#include "base/task/post_task.h"
 #include "base/trace_event/trace_config.h"
 #include "chromecast/tracing/system_tracing_common.h"
 #include "content/public/browser/browser_thread.h"
@@ -47,10 +47,9 @@ CastTracingAgent::CastTracingAgent(service_manager::Connector* connector)
                 tracing::mojom::TraceDataType::STRING,
                 false /* supports_explicit_clock_sync */,
                 base::kNullProcessId) {
-  task_runner_ =
-      base::TaskScheduler::GetInstance()->CreateSequencedTaskRunnerWithTraits(
-          {base::MayBlock(), base::TaskPriority::BEST_EFFORT,
-           base::TaskShutdownBehavior::SKIP_ON_SHUTDOWN});
+  task_runner_ = base::CreateSequencedTaskRunnerWithTraits(
+      {base::MayBlock(), base::TaskPriority::BEST_EFFORT,
+       base::TaskShutdownBehavior::SKIP_ON_SHUTDOWN});
 }
 
 CastTracingAgent::~CastTracingAgent() = default;

net/disk_cache/simple/simple_backend_impl.cc

@@ -263,11 +263,10 @@ void SimpleBackendImpl::SetWorkerPoolForTesting(
 }
 
 int SimpleBackendImpl::Init(CompletionOnceCallback completion_callback) {
-  auto worker_pool =
-      base::TaskScheduler::GetInstance()->CreateTaskRunnerWithTraits(
-          {base::MayBlock(), base::WithBaseSyncPrimitives(),
-           base::TaskPriority::USER_BLOCKING,
-           base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN});
+  auto worker_pool = base::CreateTaskRunnerWithTraits(
+      {base::MayBlock(), base::WithBaseSyncPrimitives(),
+       base::TaskPriority::USER_BLOCKING,
+       base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN});
 
   prioritized_task_runner_ =
       base::MakeRefCounted<net::PrioritizedTaskRunner>(worker_pool);