Register a thread affine TaskExecutor for ThreadPool tasks

Needed for PostTaskHere and base::CurrentThread trait. This has the same
lifetime as SequencedTaskRunnerHandle & ThreadTaskRunnerHandle.

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

base/task/post_task_unittest.cc

@@ -5,9 +5,11 @@
 #include "base/task/post_task.h"
 
 #include "base/bind_helpers.h"
+#include "base/run_loop.h"
 #include "base/task/scoped_set_task_priority_for_current_thread.h"
 #include "base/task/task_executor.h"
 #include "base/task/test_task_traits_extension.h"
+#include "base/test/bind_test_util.h"
 #include "base/test/gtest_util.h"
 #include "base/test/task_environment.h"
 #include "base/test/test_simple_task_runner.h"
@@ -17,6 +19,8 @@
 
 using ::testing::_;
 using ::testing::Invoke;
+using ::testing::IsNull;
+using ::testing::NotNull;
 using ::testing::Return;
 
 namespace base {
@@ -177,6 +181,50 @@ TEST_F(PostTaskTestWithExecutor, PostTaskToTaskExecutor) {
   }
 }
 
+TEST_F(PostTaskTestWithExecutor,
+       ThreadPoolTaskRunnerGetTaskExecutorForCurrentThread) {
+  auto task_runner = CreateTaskRunner({ThreadPool()});
+  RunLoop run_loop;
+
+  EXPECT_TRUE(task_runner->PostTask(
+      FROM_HERE, BindLambdaForTesting([&]() {
+        // We don't have an executor for a ThreadPool task runner becuse they
+        // are for one shot tasks.
+        EXPECT_THAT(GetTaskExecutorForCurrentThread(), IsNull());
+        run_loop.Quit();
+      })));
+
+  run_loop.Run();
+}
+
+TEST_F(PostTaskTestWithExecutor,
+       ThreadPoolSequencedTaskRunnerGetTaskExecutorForCurrentThread) {
+  auto sequenced_task_runner = CreateSequencedTaskRunner({ThreadPool()});
+  RunLoop run_loop;
+
+  EXPECT_TRUE(sequenced_task_runner->PostTask(
+      FROM_HERE, BindLambdaForTesting([&]() {
+        EXPECT_THAT(GetTaskExecutorForCurrentThread(), NotNull());
+        run_loop.Quit();
+      })));
+
+  run_loop.Run();
+}
+
+TEST_F(PostTaskTestWithExecutor,
+       ThreadPoolSingleThreadTaskRunnerGetTaskExecutorForCurrentThread) {
+  auto single_thread_task_runner = CreateSingleThreadTaskRunner({ThreadPool()});
+  RunLoop run_loop;
+
+  EXPECT_TRUE(single_thread_task_runner->PostTask(
+      FROM_HERE, BindLambdaForTesting([&]() {
+        EXPECT_THAT(GetTaskExecutorForCurrentThread(), NotNull());
+        run_loop.Quit();
+      })));
+
+  run_loop.Run();
+}
+
 TEST_F(PostTaskTestWithExecutor, RegisterExecutorTwice) {
   testing::FLAGS_gtest_death_test_style = "threadsafe";
   EXPECT_DCHECK_DEATH(

base/task/thread_pool/task_tracker.cc

@@ -20,6 +20,7 @@
 #include "base/sequence_token.h"
 #include "base/synchronization/condition_variable.h"
 #include "base/task/scoped_set_task_priority_for_current_thread.h"
+#include "base/task/task_executor.h"
 #include "base/threading/sequence_local_storage_map.h"
 #include "base/threading/sequenced_task_runner_handle.h"
 #include "base/threading/thread_restrictions.h"
@@ -27,6 +28,7 @@
 #include "base/time/time.h"
 #include "base/trace_event/trace_event.h"
 #include "base/values.h"
+#include "build/build_config.h"
 
 namespace base {
 namespace internal {
@@ -111,6 +113,61 @@ bool HasLogBestEffortTasksSwitch() {
              switches::kLogBestEffortTasks);
 }
 
+// Needed for PostTaskHere and CurrentThread. This executor lives for the
+// duration of a threadpool task invocation.
+class EphemeralTaskExecutor : public TaskExecutor {
+ public:
+  // |sequenced_task_runner| and |single_thread_task_runner| must outlive this
+  // EphemeralTaskExecutor.
+  EphemeralTaskExecutor(SequencedTaskRunner* sequenced_task_runner,
+                        SingleThreadTaskRunner* single_thread_task_runner)
+      : sequenced_task_runner_(sequenced_task_runner),
+        single_thread_task_runner_(single_thread_task_runner) {
+    SetTaskExecutorForCurrentThread(this);
+  }
+
+  ~EphemeralTaskExecutor() override {
+    SetTaskExecutorForCurrentThread(nullptr);
+  }
+
+  // TaskExecutor:
+  bool PostDelayedTask(const Location& from_here,
+                       const TaskTraits& traits,
+                       OnceClosure task,
+                       TimeDelta delay) override {
+    return sequenced_task_runner_->PostDelayedTask(from_here, std::move(task),
+                                                   delay);
+  }
+
+  scoped_refptr<TaskRunner> CreateTaskRunner(
+      const TaskTraits& traits) override {
+    return sequenced_task_runner_;
+  }
+
+  scoped_refptr<SequencedTaskRunner> CreateSequencedTaskRunner(
+      const TaskTraits& traits) override {
+    return sequenced_task_runner_;
+  }
+
+  scoped_refptr<SingleThreadTaskRunner> CreateSingleThreadTaskRunner(
+      const TaskTraits& traits,
+      SingleThreadTaskRunnerThreadMode thread_mode) override {
+    return single_thread_task_runner_;
+  }
+
+#if defined(OS_WIN)
+  scoped_refptr<SingleThreadTaskRunner> CreateCOMSTATaskRunner(
+      const TaskTraits& traits,
+      SingleThreadTaskRunnerThreadMode thread_mode) override {
+    return single_thread_task_runner_;
+  }
+#endif  // defined(OS_WIN)
+
+ private:
+  SequencedTaskRunner* sequenced_task_runner_;
+  SingleThreadTaskRunner* single_thread_task_runner_;
+};
+
 }  // namespace
 
 // Atomic internal state used by TaskTracker to track items that are blocking
@@ -479,6 +536,7 @@ void TaskTracker::RunTask(Task task,
     // Set up TaskRunnerHandle as expected for the scope of the task.
     Optional<SequencedTaskRunnerHandle> sequenced_task_runner_handle;
     Optional<ThreadTaskRunnerHandle> single_thread_task_runner_handle;
+    Optional<EphemeralTaskExecutor> ephemiral_task_executor;
     switch (task_source->execution_mode()) {
       case TaskSourceExecutionMode::kJob:
       case TaskSourceExecutionMode::kParallel:
@@ -487,11 +545,17 @@ void TaskTracker::RunTask(Task task,
         DCHECK(task_source->task_runner());
         sequenced_task_runner_handle.emplace(
             static_cast<SequencedTaskRunner*>(task_source->task_runner()));
+        ephemiral_task_executor.emplace(
+            static_cast<SequencedTaskRunner*>(task_source->task_runner()),
+            nullptr);
         break;
       case TaskSourceExecutionMode::kSingleThread:
         DCHECK(task_source->task_runner());
         single_thread_task_runner_handle.emplace(
             static_cast<SingleThreadTaskRunner*>(task_source->task_runner()));
+        ephemiral_task_executor.emplace(
+            static_cast<SequencedTaskRunner*>(task_source->task_runner()),
+            static_cast<SingleThreadTaskRunner*>(task_source->task_runner()));
         break;
     }