Register a TaskExecutor for base::Thread and base::SingleThreadTaskExecutor

This is needed for implementing PostTaskHere and base::CurrentThread traits.
In a follow on patch SimpleTaskExecutor will be added for blink threads.

Bug: 835323
Change-Id: Ice25d2a93fc7a6a10c0a4d4e43aa82a97462e3c9
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1832280
Commit-Queue: Alex Clarke <alexclarke@chromium.org>
Reviewed-by: François Doray <fdoray@chromium.org>
Reviewed-by: Sami Kyöstilä <skyostil@chromium.org>
Cr-Commit-Position: refs/heads/master@{#702046}

base/BUILD.gn

@@ -815,6 +815,8 @@ jumbo_component("base") {
     "task/sequence_manager/work_queue.h",
     "task/sequence_manager/work_queue_sets.cc",
     "task/sequence_manager/work_queue_sets.h",
+    "task/simple_task_executor.cc",
+    "task/simple_task_executor.h",
     "task/single_thread_task_executor.cc",
     "task/single_thread_task_executor.h",
     "task/single_thread_task_runner_thread_mode.h",
@@ -2711,6 +2713,7 @@ test("base_unittests") {
     "task/sequence_manager/work_deduplicator_unittest.cc",
     "task/sequence_manager/work_queue_sets_unittest.cc",
     "task/sequence_manager/work_queue_unittest.cc",
+    "task/single_thread_task_executor_unittest.cc",
     "task/task_traits_extension_unittest.cc",
     "task/task_traits_unittest.cc",
     "task/thread_pool/can_run_policy_test.h",

base/task/simple_task_executor.cc

+// Copyright 2019 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/simple_task_executor.h"
+
+namespace base {
+
+SimpleTaskExecutor::SimpleTaskExecutor(
+    scoped_refptr<SingleThreadTaskRunner> task_queue)
+    : task_queue_(std::move(task_queue)),
+      previous_task_executor_(GetTaskExecutorForCurrentThread()) {
+  DCHECK(task_queue_);
+  // The TaskExecutor API does not expect nesting, but this can happen in tests
+  // so we have to work around it here.
+  if (previous_task_executor_)
+    SetTaskExecutorForCurrentThread(nullptr);
+  SetTaskExecutorForCurrentThread(this);
+}
+
+SimpleTaskExecutor::~SimpleTaskExecutor() {
+  if (previous_task_executor_)
+    SetTaskExecutorForCurrentThread(nullptr);
+  SetTaskExecutorForCurrentThread(previous_task_executor_);
+}
+
+bool SimpleTaskExecutor::PostDelayedTask(const Location& from_here,
+                                         const TaskTraits& traits,
+                                         OnceClosure task,
+                                         TimeDelta delay) {
+  return task_queue_->PostDelayedTask(from_here, std::move(task), delay);
+}
+
+scoped_refptr<TaskRunner> SimpleTaskExecutor::CreateTaskRunner(
+    const TaskTraits& traits) {
+  return task_queue_;
+}
+
+scoped_refptr<SequencedTaskRunner>
+SimpleTaskExecutor::CreateSequencedTaskRunner(const TaskTraits& traits) {
+  return task_queue_;
+}
+
+scoped_refptr<SingleThreadTaskRunner>
+SimpleTaskExecutor::CreateSingleThreadTaskRunner(
+    const TaskTraits& traits,
+    SingleThreadTaskRunnerThreadMode thread_mode) {
+  return task_queue_;
+}
+
+#if defined(OS_WIN)
+scoped_refptr<SingleThreadTaskRunner>
+SimpleTaskExecutor::CreateCOMSTATaskRunner(
+    const TaskTraits& traits,
+    SingleThreadTaskRunnerThreadMode thread_mode) {
+  // It seems pretty unlikely this will be used on a comsta task thread.
+  NOTREACHED();
+  return task_queue_;
+}
+#endif  // defined(OS_WIN)
+
+}  // namespace base

base/task/simple_task_executor.h

+// Copyright 2019 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_SIMPLE_TASK_EXECUTOR_H_
+#define BASE_TASK_SIMPLE_TASK_EXECUTOR_H_
+
+#include "base/task/task_executor.h"
+#include "build/build_config.h"
+
+namespace base {
+
+// A simple TaskExecutor with exactly one SingleThreadTaskRunner.
+// Must be instantiated and destroyed on the thread that runs tasks for the
+// SingleThreadTaskRunner.
+class BASE_EXPORT SimpleTaskExecutor : public TaskExecutor {
+ public:
+  explicit SimpleTaskExecutor(scoped_refptr<SingleThreadTaskRunner> task_queue);
+
+  ~SimpleTaskExecutor() override;
+
+  bool PostDelayedTask(const Location& from_here,
+                       const TaskTraits& traits,
+                       OnceClosure task,
+                       TimeDelta delay) override;
+
+  scoped_refptr<TaskRunner> CreateTaskRunner(const TaskTraits& traits) override;
+
+  scoped_refptr<SequencedTaskRunner> CreateSequencedTaskRunner(
+      const TaskTraits& traits) override;
+
+  scoped_refptr<SingleThreadTaskRunner> CreateSingleThreadTaskRunner(
+      const TaskTraits& traits,
+      SingleThreadTaskRunnerThreadMode thread_mode) override;
+
+#if defined(OS_WIN)
+  scoped_refptr<SingleThreadTaskRunner> CreateCOMSTATaskRunner(
+      const TaskTraits& traits,
+      SingleThreadTaskRunnerThreadMode thread_mode) override;
+#endif  // defined(OS_WIN)
+
+ private:
+  const scoped_refptr<SingleThreadTaskRunner> task_queue_;
+
+  // In tests there may already be a TaskExecutor registered for the thread, we
+  // keep tack of the previous TaskExecutor and restored it upon destruction.
+  TaskExecutor* const previous_task_executor_;
+};
+
+}  // namespace base
+
+#endif  // BASE_TASK_SIMPLE_TASK_EXECUTOR_H_

base/task/single_thread_task_executor.cc

@@ -17,7 +17,8 @@ SingleThreadTaskExecutor::SingleThreadTaskExecutor(MessagePumpType type)
               .Build())),
       default_task_queue_(sequence_manager_->CreateTaskQueue(
           sequence_manager::TaskQueue::Spec("default_tq"))),
-      type_(type) {
+      type_(type),
+      simple_task_executor_(task_runner()) {
   sequence_manager_->SetDefaultTaskRunner(default_task_queue_->task_runner());
   sequence_manager_->BindToMessagePump(MessagePump::Create(type));
 

base/task/single_thread_task_executor.h

@@ -11,6 +11,7 @@
 #include "base/memory/scoped_refptr.h"
 #include "base/message_loop/message_pump_type.h"
 #include "base/single_thread_task_runner.h"
+#include "base/task/simple_task_executor.h"
 
 namespace base {
 
@@ -40,6 +41,7 @@ class BASE_EXPORT SingleThreadTaskExecutor {
   std::unique_ptr<sequence_manager::SequenceManager> sequence_manager_;
   scoped_refptr<sequence_manager::TaskQueue> default_task_queue_;
   MessagePumpType type_;
+  SimpleTaskExecutor simple_task_executor_;
 
   DISALLOW_COPY_AND_ASSIGN(SingleThreadTaskExecutor);
 };

base/task/single_thread_task_executor_unittest.cc

+// Copyright 2019 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/single_thread_task_executor.h"
+
+#include "base/run_loop.h"
+#include "base/test/bind_test_util.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+using ::testing::IsNull;
+using ::testing::NotNull;
+
+namespace base {
+
+TEST(SingleThreadTaskExecutorTest, GetTaskExecutorForCurrentThread) {
+  EXPECT_THAT(GetTaskExecutorForCurrentThread(), IsNull());
+
+  {
+    SingleThreadTaskExecutor single_thread_task_executor;
+    EXPECT_THAT(GetTaskExecutorForCurrentThread(), NotNull());
+  }
+
+  EXPECT_THAT(GetTaskExecutorForCurrentThread(), IsNull());
+}
+
+TEST(SingleThreadTaskExecutorTest,
+     GetTaskExecutorForCurrentThreadInPostedTask) {
+  SingleThreadTaskExecutor single_thread_task_executor;
+  TaskExecutor* task_executor = GetTaskExecutorForCurrentThread();
+
+  EXPECT_THAT(task_executor, NotNull());
+
+  RunLoop run_loop;
+  single_thread_task_executor.task_runner()->PostTask(
+      FROM_HERE, BindLambdaForTesting([&]() {
+        EXPECT_EQ(GetTaskExecutorForCurrentThread(), task_executor);
+        run_loop.Quit();
+      }));
+
+  run_loop.Run();
+}
+
+}  // namespace base

base/threading/thread.cc

@@ -19,6 +19,7 @@
 #include "base/synchronization/waitable_event.h"
 #include "base/task/sequence_manager/sequence_manager_impl.h"
 #include "base/task/sequence_manager/task_queue.h"
+#include "base/task/simple_task_executor.h"
 #include "base/third_party/dynamic_annotations/dynamic_annotations.h"
 #include "base/threading/thread_id_name_manager.h"
 #include "base/threading/thread_local.h"
@@ -86,6 +87,7 @@ class SequenceManagerThreadDelegate : public Thread::Delegate {
     sequence_manager_->BindToMessagePump(
         std::move(message_pump_factory_).Run());
     sequence_manager_->SetTimerSlack(timer_slack);
+    simple_task_executor_.emplace(GetDefaultTaskRunner());
   }
 
  private:
@@ -93,6 +95,7 @@ class SequenceManagerThreadDelegate : public Thread::Delegate {
       sequence_manager_;
   scoped_refptr<sequence_manager::TaskQueue> default_task_queue_;
   OnceCallback<std::unique_ptr<MessagePump>()> message_pump_factory_;
+  base::Optional<SimpleTaskExecutor> simple_task_executor_;
 };
 
 }  // namespace

base/threading/thread_unittest.cc

@@ -18,15 +18,19 @@
 #include "base/single_thread_task_runner.h"
 #include "base/synchronization/waitable_event.h"
 #include "base/task/sequence_manager/sequence_manager_impl.h"
+#include "base/task/task_executor.h"
+#include "base/test/bind_test_util.h"
 #include "base/test/gtest_util.h"
 #include "base/third_party/dynamic_annotations/dynamic_annotations.h"
 #include "base/threading/platform_thread.h"
 #include "base/time/time.h"
 #include "build/build_config.h"
+#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "testing/platform_test.h"
 
 using base::Thread;
+using ::testing::NotNull;
 
 typedef PlatformTest ThreadTest;
 
@@ -522,6 +526,23 @@ TEST_F(ThreadTest, FlushForTesting) {
   a.FlushForTesting();
 }
 
+TEST_F(ThreadTest, GetTaskExecutorForCurrentThread) {
+  Thread a("GetTaskExecutorForCurrentThread");
+  ASSERT_TRUE(a.Start());
+
+  base::WaitableEvent event(base::WaitableEvent::ResetPolicy::MANUAL,
+                            base::WaitableEvent::InitialState::NOT_SIGNALED);
+
+  a.task_runner()->PostTask(
+      FROM_HERE, base::BindLambdaForTesting([&]() {
+        EXPECT_THAT(base::GetTaskExecutorForCurrentThread(), NotNull());
+        event.Signal();
+      }));
+
+  event.Wait();
+  a.Stop();
+}
+
 namespace {
 
 class SequenceManagerThreadDelegate : public Thread::Delegate {