TaskScheduler: Set the current SequenceToken when a Task runs.

base/task_scheduler/task_tracker_unittest.cc

Index: base/task_scheduler/task_tracker_unittest.cc
diff --git a/base/task_scheduler/task_tracker_unittest.cc b/base/task_scheduler/task_tracker_unittest.cc
index 326f7aa02bc79602b34d665a2762e6e7965b7cf4..603bf73cd7fa26645a0677452ea7189ad363f18d 100644
--- a/base/task_scheduler/task_tracker_unittest.cc
+++ b/base/task_scheduler/task_tracker_unittest.cc
@@ -14,10 +14,12 @@
 #include "base/macros.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/ref_counted.h"
+#include "base/sequence_token.h"
 #include "base/sequenced_task_runner.h"
 #include "base/single_thread_task_runner.h"
 #include "base/synchronization/waitable_event.h"
 #include "base/task_scheduler/scheduler_lock.h"
+#include "base/task_scheduler/sequence.h"
 #include "base/task_scheduler/task.h"
 #include "base/task_scheduler/task_traits.h"
 #include "base/task_scheduler/test_utils.h"
@@ -36,6 +38,12 @@ namespace {
 
 constexpr size_t kLoadTestNumIterations = 100;
 
+scoped_refptr<Sequence> CreateSequenceWithTask(std::unique_ptr<Task> task) {
+  scoped_refptr<Sequence> sequence(new Sequence);
+  sequence->PushTask(std::move(task));
+  return sequence;
+}
+
 // Calls TaskTracker::Shutdown() asynchronously.
 class ThreadCallingShutdown : public SimpleThread {
  public:
@@ -69,12 +77,12 @@ class ThreadPostingAndRunningTask : public SimpleThread {
   };
 
   ThreadPostingAndRunningTask(TaskTracker* tracker,
-                              const Task* task,
+                              scoped_refptr<Sequence> sequence,
                               Action action,
                               bool expect_post_succeeds)
       : SimpleThread("ThreadPostingAndRunningTask"),
         tracker_(tracker),
-        task_(task),
+        sequence_(std::move(sequence)),
         action_(action),
         expect_post_succeeds_(expect_post_succeeds) {}
 
@@ -82,17 +90,17 @@ class ThreadPostingAndRunningTask : public SimpleThread {
   void Run() override {
     bool post_succeeded = true;
     if (action_ == Action::WILL_POST || action_ == Action::WILL_POST_AND_RUN) {
-      post_succeeded = tracker_->WillPostTask(task_);
+      post_succeeded = tracker_->WillPostTask(sequence_->PeekTask());
       EXPECT_EQ(expect_post_succeeds_, post_succeeded);
     }
     if (post_succeeded &&
         (action_ == Action::RUN || action_ == Action::WILL_POST_AND_RUN)) {
-      tracker_->RunTask(task_);
+      tracker_->RunNextTaskInSequence(sequence_.get());
     }
   }
 
   TaskTracker* const tracker_;
-  const Task* const task_;
+  const scoped_refptr<Sequence> sequence_;
   const Action action_;
   const bool expect_post_succeeds_;
 
@@ -195,7 +203,7 @@ TEST_P(TaskSchedulerTaskTrackerTest, WillPostAndRunBeforeShutdown) {
 
   // Run the task.
   EXPECT_EQ(0U, NumTasksExecuted());
-  tracker_.RunTask(task.get());
+  tracker_.RunNextTaskInSequence(CreateSequenceWithTask(std::move(task)).get());
   EXPECT_EQ(1U, NumTasksExecuted());
 
   // Shutdown() shouldn't block.
@@ -215,8 +223,8 @@ TEST_P(TaskSchedulerTaskTrackerTest, WillPostAndRunLongTaskBeforeShutdown) {
 
   // Run the task asynchronouly.
   ThreadPostingAndRunningTask thread_running_task(
-      &tracker_, blocked_task.get(), ThreadPostingAndRunningTask::Action::RUN,
-      false);
+      &tracker_, CreateSequenceWithTask(std::move(blocked_task)),
+      ThreadPostingAndRunningTask::Action::RUN, false);
   thread_running_task.Start();
 
   // Initiate shutdown while the task is running.
@@ -257,13 +265,14 @@ TEST_P(TaskSchedulerTaskTrackerTest, WillPostBeforeShutdownRunDuringShutdown) {
   // Try to run |task|. It should only run it it's BLOCK_SHUTDOWN. Otherwise it
   // should be discarded.
   EXPECT_EQ(0U, NumTasksExecuted());
-  tracker_.RunTask(task.get());
+  tracker_.RunNextTaskInSequence(CreateSequenceWithTask(std::move(task)).get());
   EXPECT_EQ(GetParam() == TaskShutdownBehavior::BLOCK_SHUTDOWN ? 1U : 0U,
             NumTasksExecuted());
   VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS();
 
   // Unblock shutdown by running the remaining BLOCK_SHUTDOWN task.
-  tracker_.RunTask(block_shutdown_task.get());
+  tracker_.RunNextTaskInSequence(
+      CreateSequenceWithTask(std::move(block_shutdown_task)).get());
   EXPECT_EQ(GetParam() == TaskShutdownBehavior::BLOCK_SHUTDOWN ? 2U : 1U,
             NumTasksExecuted());
   WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
@@ -282,7 +291,8 @@ TEST_P(TaskSchedulerTaskTrackerTest, WillPostBeforeShutdownRunAfterShutdown) {
     VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS();
 
     // Run the task to unblock shutdown.
-    tracker_.RunTask(task.get());
+    tracker_.RunNextTaskInSequence(
+        CreateSequenceWithTask(std::move(task)).get());
     EXPECT_EQ(1U, NumTasksExecuted());
     WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
 
@@ -293,7 +303,8 @@ TEST_P(TaskSchedulerTaskTrackerTest, WillPostBeforeShutdownRunAfterShutdown) {
     WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
 
     // The task shouldn't be allowed to run after shutdown.
-    tracker_.RunTask(task.get());
+    tracker_.RunNextTaskInSequence(
+        CreateSequenceWithTask(std::move(task)).get());
     EXPECT_EQ(0U, NumTasksExecuted());
   }
 }
@@ -316,7 +327,8 @@ TEST_P(TaskSchedulerTaskTrackerTest, WillPostAndRunDuringShutdown) {
 
     // Run the BLOCK_SHUTDOWN task.
     EXPECT_EQ(0U, NumTasksExecuted());
-    tracker_.RunTask(task.get());
+    tracker_.RunNextTaskInSequence(
+        CreateSequenceWithTask(std::move(task)).get());
     EXPECT_EQ(1U, NumTasksExecuted());
   } else {
     // It shouldn't be allowed to post a non BLOCK_SHUTDOWN task.
@@ -328,7 +340,8 @@ TEST_P(TaskSchedulerTaskTrackerTest, WillPostAndRunDuringShutdown) {
 
   // Unblock shutdown by running |block_shutdown_task|.
   VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS();
-  tracker_.RunTask(block_shutdown_task.get());
+  tracker_.RunNextTaskInSequence(
+      CreateSequenceWithTask(std::move(block_shutdown_task)).get());
   EXPECT_EQ(GetParam() == TaskShutdownBehavior::BLOCK_SHUTDOWN ? 2U : 1U,
             NumTasksExecuted());
   WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
@@ -354,9 +367,10 @@ TEST_P(TaskSchedulerTaskTrackerTest, SingletonAllowed) {
       (GetParam() != TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN);
 
   TaskTracker tracker;
-  Task task(FROM_HERE, Bind(&ThreadRestrictions::AssertSingletonAllowed),
-            TaskTraits().WithShutdownBehavior(GetParam()), TimeDelta());
-  EXPECT_TRUE(tracker.WillPostTask(&task));
+  std::unique_ptr<Task> task(
+      new Task(FROM_HERE, Bind(&ThreadRestrictions::AssertSingletonAllowed),
+               TaskTraits().WithShutdownBehavior(GetParam()), TimeDelta()));
+  EXPECT_TRUE(tracker.WillPostTask(task.get()));
 
   // Set the singleton allowed bit to the opposite of what it is expected to be
   // when |tracker| runs |task| to verify that |tracker| actually sets the
@@ -365,23 +379,31 @@ TEST_P(TaskSchedulerTaskTrackerTest, SingletonAllowed) {
 
   // Running the task should fail iff the task isn't allowed to use singletons.
   if (can_use_singletons) {
-    tracker.RunTask(&task);
+    tracker.RunNextTaskInSequence(
+        CreateSequenceWithTask(std::move(task)).get());
   } else {
-    EXPECT_DCHECK_DEATH({ tracker.RunTask(&task); }, "");
+    EXPECT_DCHECK_DEATH(
+        {
+          tracker.RunNextTaskInSequence(
+              CreateSequenceWithTask(std::move(task)).get());
+        },
+        "");
   }
 }
 
-static void RunTaskRunnerHandleVerificationTask(TaskTracker* tracker,
-                                         const Task* verify_task) {
+static void RunTaskRunnerHandleVerificationTask(
+    TaskTracker* tracker,
+    std::unique_ptr<Task> verify_task) {
   // Pretend |verify_task| is posted to respect TaskTracker's contract.
-  EXPECT_TRUE(tracker->WillPostTask(verify_task));
+  EXPECT_TRUE(tracker->WillPostTask(verify_task.get()));
 
   // Confirm that the test conditions are right (no TaskRunnerHandles set
   // already).
   EXPECT_FALSE(ThreadTaskRunnerHandle::IsSet());
   EXPECT_FALSE(SequencedTaskRunnerHandle::IsSet());
 
-  tracker->RunTask(verify_task);
+  tracker->RunNextTaskInSequence(
+      CreateSequenceWithTask(std::move(verify_task)).get());
 
   // TaskRunnerHandle state is reset outside of task's scope.
   EXPECT_FALSE(ThreadTaskRunnerHandle::IsSet());
@@ -400,7 +422,7 @@ TEST_P(TaskSchedulerTaskTrackerTest, TaskRunnerHandleIsNotSetOnParallel) {
       new Task(FROM_HERE, Bind(&VerifyNoTaskRunnerHandle),
                TaskTraits().WithShutdownBehavior(GetParam()), TimeDelta()));
 
-  RunTaskRunnerHandleVerificationTask(&tracker_, verify_task.get());
+  RunTaskRunnerHandleVerificationTask(&tracker_, std::move(verify_task));
 }
 
 static void VerifySequencedTaskRunnerHandle(
@@ -423,7 +445,7 @@ TEST_P(TaskSchedulerTaskTrackerTest,
                TaskTraits().WithShutdownBehavior(GetParam()), TimeDelta()));
   verify_task->sequenced_task_runner_ref = test_task_runner;
 
-  RunTaskRunnerHandleVerificationTask(&tracker_, verify_task.get());
+  RunTaskRunnerHandleVerificationTask(&tracker_, std::move(verify_task));
 }
 
 static void VerifyThreadTaskRunnerHandle(
@@ -448,7 +470,7 @@ TEST_P(TaskSchedulerTaskTrackerTest,
                TaskTraits().WithShutdownBehavior(GetParam()), TimeDelta()));
   verify_task->single_thread_task_runner_ref = test_task_runner;
 
-  RunTaskRunnerHandleVerificationTask(&tracker_, verify_task.get());
+  RunTaskRunnerHandleVerificationTask(&tracker_, std::move(verify_task));
 }
 
 INSTANTIATE_TEST_CASE_P(
@@ -464,27 +486,48 @@ INSTANTIATE_TEST_CASE_P(
     TaskSchedulerTaskTrackerTest,
     ::testing::Values(TaskShutdownBehavior::BLOCK_SHUTDOWN));
 
+namespace {
+
+void ExpectSequenceToken(SequenceToken sequence_token) {
+  EXPECT_EQ(sequence_token, SequenceToken::GetForCurrentThread());
+}
+
+}  // namespace
+
+// Verify that SequenceToken::GetForCurrentThread() returns the Sequence's token
+// when a Task runs.
+TEST_F(TaskSchedulerTaskTrackerTest, CurrentSequenceToken) {
+  scoped_refptr<Sequence> sequence(new Sequence);
+  sequence->PushTask(WrapUnique(
+      new Task(FROM_HERE, Bind(&ExpectSequenceToken, sequence->token()),
+               TaskTraits(), TimeDelta())));
+  tracker_.WillPostTask(sequence->PeekTask());
+
+  EXPECT_FALSE(SequenceToken::GetForCurrentThread().IsValid());
+  tracker_.RunNextTaskInSequence(sequence.get());
+  EXPECT_FALSE(SequenceToken::GetForCurrentThread().IsValid());
+}
+
 TEST_F(TaskSchedulerTaskTrackerTest, LoadWillPostAndRunBeforeShutdown) {
   // Post and run tasks asynchronously.
-  std::vector<std::unique_ptr<Task>> tasks;
   std::vector<std::unique_ptr<ThreadPostingAndRunningTask>> threads;
 
   for (size_t i = 0; i < kLoadTestNumIterations; ++i) {
-    tasks.push_back(CreateTask(TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN));
     threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
-        &tracker_, tasks.back().get(),
+        &tracker_, CreateSequenceWithTask(
+                       CreateTask(TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN)),
         ThreadPostingAndRunningTask::Action::WILL_POST_AND_RUN, true)));
     threads.back()->Start();
 
-    tasks.push_back(CreateTask(TaskShutdownBehavior::SKIP_ON_SHUTDOWN));
     threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
-        &tracker_, tasks.back().get(),
+        &tracker_, CreateSequenceWithTask(
+                       CreateTask(TaskShutdownBehavior::SKIP_ON_SHUTDOWN)),
         ThreadPostingAndRunningTask::Action::WILL_POST_AND_RUN, true)));
     threads.back()->Start();
 
-    tasks.push_back(CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN));
     threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
-        &tracker_, tasks.back().get(),
+        &tracker_, CreateSequenceWithTask(
+                       CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN)),
         ThreadPostingAndRunningTask::Action::WILL_POST_AND_RUN, true)));
     threads.back()->Start();
   }
@@ -502,25 +545,28 @@ TEST_F(TaskSchedulerTaskTrackerTest, LoadWillPostAndRunBeforeShutdown) {
 TEST_F(TaskSchedulerTaskTrackerTest,
        LoadWillPostBeforeShutdownAndRunDuringShutdown) {
   // Post tasks asynchronously.
-  std::vector<std::unique_ptr<Task>> tasks;
+  std::vector<scoped_refptr<Sequence>> sequences;
   std::vector<std::unique_ptr<ThreadPostingAndRunningTask>> post_threads;
 
   for (size_t i = 0; i < kLoadTestNumIterations; ++i) {
-    tasks.push_back(CreateTask(TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN));
+    sequences.push_back(CreateSequenceWithTask(
+        CreateTask(TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN)));
     post_threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
-        &tracker_, tasks.back().get(),
+        &tracker_, sequences.back(),
         ThreadPostingAndRunningTask::Action::WILL_POST, true)));
     post_threads.back()->Start();
 
-    tasks.push_back(CreateTask(TaskShutdownBehavior::SKIP_ON_SHUTDOWN));
+    sequences.push_back(CreateSequenceWithTask(
+        CreateTask(TaskShutdownBehavior::SKIP_ON_SHUTDOWN)));
     post_threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
-        &tracker_, tasks.back().get(),
+        &tracker_, sequences.back(),
         ThreadPostingAndRunningTask::Action::WILL_POST, true)));
     post_threads.back()->Start();
 
-    tasks.push_back(CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN));
+    sequences.push_back(CreateSequenceWithTask(
+        CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN)));
     post_threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
-        &tracker_, tasks.back().get(),
+        &tracker_, sequences.back(),
         ThreadPostingAndRunningTask::Action::WILL_POST, true)));
     post_threads.back()->Start();
   }
@@ -534,10 +580,9 @@ TEST_F(TaskSchedulerTaskTrackerTest,
   // Run tasks asynchronously.
   std::vector<std::unique_ptr<ThreadPostingAndRunningTask>> run_threads;
 
-  for (const auto& task : tasks) {
+  for (const auto& sequence : sequences) {
     run_threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
-        &tracker_, task.get(), ThreadPostingAndRunningTask::Action::RUN,
-        false)));
+        &tracker_, sequence, ThreadPostingAndRunningTask::Action::RUN, false)));
     run_threads.back()->Start();
   }
 
@@ -561,25 +606,24 @@ TEST_F(TaskSchedulerTaskTrackerTest, LoadWillPostAndRunDuringShutdown) {
   CallShutdownAsync();
 
   // Post and run tasks asynchronously.
-  std::vector<std::unique_ptr<Task>> tasks;
   std::vector<std::unique_ptr<ThreadPostingAndRunningTask>> threads;
 
   for (size_t i = 0; i < kLoadTestNumIterations; ++i) {
-    tasks.push_back(CreateTask(TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN));
     threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
-        &tracker_, tasks.back().get(),
+        &tracker_, CreateSequenceWithTask(
+                       CreateTask(TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN)),
         ThreadPostingAndRunningTask::Action::WILL_POST_AND_RUN, false)));
     threads.back()->Start();
 
-    tasks.push_back(CreateTask(TaskShutdownBehavior::SKIP_ON_SHUTDOWN));
     threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
-        &tracker_, tasks.back().get(),
+        &tracker_, CreateSequenceWithTask(
+                       CreateTask(TaskShutdownBehavior::SKIP_ON_SHUTDOWN)),
         ThreadPostingAndRunningTask::Action::WILL_POST_AND_RUN, false)));
     threads.back()->Start();
 
-    tasks.push_back(CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN));
     threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
-        &tracker_, tasks.back().get(),
+        &tracker_, CreateSequenceWithTask(
+                       CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN)),
         ThreadPostingAndRunningTask::Action::WILL_POST_AND_RUN, true)));
     threads.back()->Start();
   }
@@ -594,7 +638,8 @@ TEST_F(TaskSchedulerTaskTrackerTest, LoadWillPostAndRunDuringShutdown) {
   VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS();
 
   // Unblock shutdown by running |block_shutdown_task|.
-  tracker_.RunTask(block_shutdown_task.get());
+  tracker_.RunNextTaskInSequence(
+      CreateSequenceWithTask(std::move(block_shutdown_task)).get());
   EXPECT_EQ(kLoadTestNumIterations + 1, NumTasksExecuted());
   WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
 }