TaskScheduler: Add TaskRunnerHandle support to TaskScheduler tasks

base/task_scheduler/task_tracker.cc

 // Copyright 2016 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_scheduler/task_tracker.h"
 
 #include "base/callback.h"
 #include "base/debug/task_annotator.h"
 #include "base/metrics/histogram_macros.h"
+#include "base/thread_task_runner_handle.h"
+#include "base/threading/sequenced_task_runner_handle.h"
 
 namespace base {
 namespace internal {
 
 namespace {
 
 const char kQueueFunctionName[] = "base::PostTask";
 
 // Upper bound for the
 // TaskScheduler.BlockShutdownTasksPostedDuringShutdown histogram.
@@ skipping 51 matching lines @@
 }
 
 void TaskTracker::RunTask(const Task* task) {
   DCHECK(task);
 
   const TaskShutdownBehavior shutdown_behavior =
       task->traits.shutdown_behavior();
   if (!BeforeRunTask(shutdown_behavior))
     return;
 
-  debug::TaskAnnotator task_annotator;
-  task_annotator.RunTask(kQueueFunctionName, *task);
+  {
+    // Set up TaskRunnerHandle as expected for the scope of the task.
+    std::unique_ptr<SequencedTaskRunnerHandle> sequenced_trh;
+    std::unique_ptr<ThreadTaskRunnerHandle> single_thread_trh;
+    DCHECK(!task->sequenced_task_runner_ref ||
+           !task->single_thread_task_runner_ref);
+    if (task->sequenced_task_runner_ref) {
+      sequenced_trh.reset(
+          new SequencedTaskRunnerHandle(task->sequenced_task_runner_ref));
+    } else if (task->single_thread_task_runner_ref) {
+      single_thread_trh.reset(
+          new ThreadTaskRunnerHandle(task->single_thread_task_runner_ref));
+    }
+
+    debug::TaskAnnotator task_annotator;
+    task_annotator.RunTask(kQueueFunctionName, *task);
+  }
 
   AfterRunTask(shutdown_behavior);
 }
 
 bool TaskTracker::IsShuttingDownForTesting() const {
   AutoSchedulerLock auto_lock(lock_);
   return !!shutdown_cv_;
 }
 
 bool TaskTracker::BeforePostTask(TaskShutdownBehavior shutdown_behavior) {
@@ skipping 78 matching lines @@
     AutoSchedulerLock auto_lock(lock_);
     DCHECK_GT(num_tasks_blocking_shutdown_, 0U);
     --num_tasks_blocking_shutdown_;
     if (num_tasks_blocking_shutdown_ == 0 && shutdown_cv_)
       shutdown_cv_->Signal();
   }
 }
 
 }  // namespace internal
 }  // namespace base