Use the saved task runner if TestTaskScheduler is re-entered from a task executed by RunTask.

base/test/scoped_task_scheduler.cc

Index: base/test/scoped_task_scheduler.cc
diff --git a/base/test/scoped_task_scheduler.cc b/base/test/scoped_task_scheduler.cc
index 7bcd88adc0ab3c8eb95da13b6e17e982d5ebb3e5..96cd3ce5eae5df9718cb3aa9c0904faf20ffa244 100644
--- a/base/test/scoped_task_scheduler.cc
+++ b/base/test/scoped_task_scheduler.cc
@@ -70,6 +70,14 @@ class TestTaskScheduler : public TaskScheduler {
   bool RunsTasksOnCurrentThread() const;
 
  private:
+  // Returns |saved_task_runner_| if called from inside RunTask.

[fdoray, 2017/02/22 16:23:52]

Returns the TaskRunner to which this TaskScheduler forwards tasks. It may be
|message_loop_->task_runner()| or a reference to it saved on entry to RunTask().

[Joe Mason, 2017/02/22 20:23:30]

Done.

+  scoped_refptr<SingleThreadTaskRunner> CurrentTaskRunner() const {

[fdoray, 2017/02/22 16:23:52]

s/CurrentTaskRunner/MessageLoopTaskRunner/

"Current" make it sound like it changes often, which isn't the case. This will
return the same pointer, obtained from a different source, when called outside
or inside of RunTask().

[Joe Mason, 2017/02/22 20:23:30]

Done.

+    if (saved_task_runner_)
+      return saved_task_runner_;
+    DCHECK(message_loop_->task_runner());
+    return message_loop_->task_runner();
+  }
+
   // |message_loop_owned_| will be non-null if this TestTaskScheduler owns the
   // MessageLoop (wasn't provided an external one at construction).
   // |message_loop_| will always be set and is used by this TestTaskScheduler to
@@ -77,6 +85,11 @@ class TestTaskScheduler : public TaskScheduler {
   std::unique_ptr<MessageLoop> message_loop_owned_;
   MessageLoop* message_loop_;
 
+  // A copy of |message_loop_|'s task_runner that is saved on entry to RunTask.

[fdoray, 2017/02/22 16:23:51]

A reference to |message_loop_->task_runner()| saved on entry to RunTask(). This
is required because RunTask() overrides |message_loop_->task_runner()|.

Note: |message_loop_->task_runner()| is accessed directly outside of RunTask()
to guarantee that ScopedTaskScheduler always uses the latest TaskRunner set by
external code.

[Joe Mason, 2017/02/22 20:23:30]

Done.

+  // It will be used if the task posted by RunTask calls any TestTaskScheduler
+  // methods.
+  scoped_refptr<SingleThreadTaskRunner> saved_task_runner_;
+
   // Handles shutdown behaviors and sets up the environment to run a task.
   internal::TaskTracker task_tracker_;
 
@@ -182,7 +195,7 @@ bool TestTaskScheduler::PostTask(std::unique_ptr<internal::Task> task,
   if (!task_tracker_.WillPostTask(task.get()))
     return false;
   internal::Task* const task_ptr = task.get();
-  return message_loop_->task_runner()->PostDelayedTask(
+  return CurrentTaskRunner()->PostDelayedTask(
       task_ptr->posted_from, Bind(&TestTaskScheduler::RunTask, Unretained(this),
                                   Passed(&task), sequence_token),
       task_ptr->delay);
@@ -192,8 +205,8 @@ void TestTaskScheduler::RunTask(std::unique_ptr<internal::Task> task,
                                 const SequenceToken& sequence_token) {
   // Clear the MessageLoop TaskRunner to allow TaskTracker to register its own
   // Thread/SequencedTaskRunnerHandle as appropriate.
-  scoped_refptr<SingleThreadTaskRunner> saved_task_runner =
-      MessageLoop::current()->task_runner();
+  DCHECK(!saved_task_runner_);

[fdoray, 2017/02/22 16:23:51]

Move these 2 lines before the comment. Add separate comment.

[Joe Mason, 2017/02/22 20:23:30]

Done.

I don't think it needs a comment because assigning to "saved_task_runner_" is
self-documenting.

+  saved_task_runner_ = MessageLoop::current()->task_runner();
   MessageLoop::current()->ClearTaskRunnerForTesting();
 
   // Run the task.
@@ -202,11 +215,12 @@ void TestTaskScheduler::RunTask(std::unique_ptr<internal::Task> task,
                                              : SequenceToken::Create());
 
   // Restore the MessageLoop TaskRunner.

[fdoray, 2017/02/22 16:23:52]

DCHECK(!MessageLoop::current()->task_runner());

[Joe Mason, 2017/02/22 20:23:30]

Done.

-  MessageLoop::current()->SetTaskRunner(saved_task_runner);
+  MessageLoop::current()->SetTaskRunner(saved_task_runner_);
+  saved_task_runner_ = nullptr;
 }
 
 bool TestTaskScheduler::RunsTasksOnCurrentThread() const {
-  return message_loop_->task_runner()->RunsTasksOnCurrentThread();
+  return CurrentTaskRunner()->RunsTasksOnCurrentThread();
 }
 
 TestTaskSchedulerTaskRunner::TestTaskSchedulerTaskRunner(