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..0976a4cbbc263d620e3559c4b09dafd30fbb9db0 100644
--- a/base/threading/sequenced_task_runner_handle.cc
+++ b/base/threading/sequenced_task_runner_handle.cc
@@ -4,28 +4,60 @@
 
 #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)
+    return handle->task_runner_;
+
+  // Return the SequencedTaskRunner obtained from SequencedWorkerPool, if any.

[danakj, 2016/04/29 22:36:34]

These can't combine right? Maybe DCHECK this is null in the return block above?

Can SequencedWorkerPool use this new mechanism to bind itself to the thread?

[gab, 2016/05/02 14:36:54]

Good point, DCHECKs added.
SequencedWorkerPool doesn't have a reference to the TaskRunner that posted the
task so it actually creates a new one on demand (ref.
SequencedWorkerPool::GetSequencedTaskRunnerForCurrentThread()). Therefore I
don't think we want to have it preemptively create one for every task just in
case.
And any other solution to keep a ref to the TaskRunner would be overly involved
IMO given it's soon being deprecated by TaskScheduler anyways :-).

[danakj, 2016/05/02 20:12:13]

Ah, right :/

   scoped_refptr<base::SequencedTaskRunner> task_runner =
       SequencedWorkerPool::GetSequencedTaskRunnerForCurrentThread();
   if (task_runner)
     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