TaskScheduler: Add TaskRunnerHandle support to TaskScheduler tasks

base/threading/sequenced_task_runner_handle.cc

Index: base/threading/sequenced_task_runner_handle.cc
diff --git a/base/threading/sequenced_task_runner_handle.cc b/base/threading/sequenced_task_runner_handle.cc
index 081f11f78b31eee9ecb1c6143e61f30a7bbee9c9..6bc75fac88143cf9088fa177e7f032383b63a7d7 100644
--- a/base/threading/sequenced_task_runner_handle.cc
+++ b/base/threading/sequenced_task_runner_handle.cc
@@ -4,28 +4,66 @@
 
 #include "base/threading/sequenced_task_runner_handle.h"
 
+#include <utility>
+
+#include "base/lazy_instance.h"
+#include "base/logging.h"
 #include "base/sequenced_task_runner.h"
 #include "base/thread_task_runner_handle.h"
 #include "base/threading/sequenced_worker_pool.h"
+#include "base/threading/thread_local.h"
 
 namespace base {
 
+namespace {
+
+base::LazyInstance<base::ThreadLocalPointer<SequencedTaskRunnerHandle>>::Leaky
+    lazy_tls_ptr = LAZY_INSTANCE_INITIALIZER;
+
+}  // namespace
+
 // static
 scoped_refptr<SequencedTaskRunner> SequencedTaskRunnerHandle::Get() {
-  // Return the SequencedTaskRunner if found or the SingleThreadedTaskRunner for
-  // the current thread otherwise.
+  // Return the registered SequencedTaskRunner, if any.
+  const SequencedTaskRunnerHandle* handle = lazy_tls_ptr.Pointer()->Get();
+  if (handle) {
+    // Various modes of setting SequencedTaskRunnerHandle don't combine.
+    DCHECK(!base::ThreadTaskRunnerHandle::IsSet());
+    DCHECK(!SequencedWorkerPool::GetSequencedTaskRunnerForCurrentThread());
+    return handle->task_runner_;
+  }
+
+  // Return the SequencedTaskRunner obtained from SequencedWorkerPool, if any.
   scoped_refptr<base::SequencedTaskRunner> task_runner =
       SequencedWorkerPool::GetSequencedTaskRunnerForCurrentThread();
-  if (task_runner)
+  if (task_runner) {
+    DCHECK(!base::ThreadTaskRunnerHandle::IsSet());
     return task_runner;
+  }
 
+  // Return the SingleThreadTaskRunner for the current thread otherwise.
   return base::ThreadTaskRunnerHandle::Get();
 }
 
 // static
 bool SequencedTaskRunnerHandle::IsSet() {
-  return SequencedWorkerPool::GetWorkerPoolForCurrentThread() ||
+  return lazy_tls_ptr.Pointer()->Get() ||
+         SequencedWorkerPool::GetWorkerPoolForCurrentThread() ||
          base::ThreadTaskRunnerHandle::IsSet();
 }
 
+SequencedTaskRunnerHandle::SequencedTaskRunnerHandle(
+    scoped_refptr<SequencedTaskRunner> task_runner)
+    : task_runner_(std::move(task_runner)) {
+  DCHECK(task_runner_->RunsTasksOnCurrentThread());
+  DCHECK(!SequencedTaskRunnerHandle::IsSet());
+  lazy_tls_ptr.Pointer()->Set(this);
+}
+
+SequencedTaskRunnerHandle::~SequencedTaskRunnerHandle() {
+  DCHECK(task_runner_->RunsTasksOnCurrentThread());
+  DCHECK_EQ(lazy_tls_ptr.Pointer()->Get(), this);
+  lazy_tls_ptr.Pointer()->Set(nullptr);
+}
+
 }  // namespace base