[scheduler] Registering the current task executor.

I'm saving to TLS the current task executor on registration, and providing
it for the base::CurrentThread() trait.

Bug: 968047
Change-Id: I753213ff57249ed55d10cb0a576ec56975a4088e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1642562Reviewed-by: Alex Clarke <alexclarke@chromium.org>
Reviewed-by: Karolina Soltys <ksolt@chromium.org>
Commit-Queue: Karolina Soltys <ksolt@chromium.org>
Cr-Commit-Position: refs/heads/master@{#666291}

base/task/task_executor.cc

@@ -6,8 +6,10 @@
 
 #include <type_traits>
 
+#include "base/no_destructor.h"
 #include "base/task/task_traits.h"
 #include "base/task/task_traits_extension.h"
+#include "base/threading/thread_local.h"
 
 namespace base {
 
@@ -28,12 +30,18 @@ static_assert(
     TaskTraitsExtensionStorage::kInvalidExtensionId == 0,
     "TaskExecutorMap depends on 0 being an invalid TaskTraits extension ID");
 
+ThreadLocalPointer<TaskExecutor>* GetTLSForCurrentTaskExecutor() {
+  static NoDestructor<ThreadLocalPointer<TaskExecutor>> instance;
+  return instance.get();
+}
+
 }  // 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);
+  GetTLSForCurrentTaskExecutor()->Set(task_executor);
 
   (*GetTaskExecutorMap())[extension_id - 1] = task_executor;
 }
@@ -42,11 +50,16 @@ 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) {
+  if (traits.use_current_thread()) {
+    auto* current_task_executor = GetTLSForCurrentTaskExecutor()->Get();
+    DCHECK(current_task_executor);
+    return current_task_executor;
+  }
+
   uint8_t extension_id = traits.extension_id();
   if (extension_id != TaskTraitsExtensionStorage::kInvalidExtensionId) {
     TaskExecutor* executor = (*GetTaskExecutorMap())[extension_id - 1];