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

base/test/scoped_task_scheduler.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/test/scoped_task_scheduler.h"
 
 #include <memory>
 #include <utility>
 #include <vector>
 
@@ skipping 52 matching lines @@
                 const SequenceToken& sequence_token);
 
   // Runs |task| with |sequence_token| using this TaskScheduler's TaskTracker.
   void RunTask(std::unique_ptr<internal::Task> task,
                const SequenceToken& sequence_token);
 
   // Returns true if this TaskScheduler runs its tasks on the current thread.
   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
   // run tasks.
   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_;
 
   DISALLOW_COPY_AND_ASSIGN(TestTaskScheduler);
 };
 
 class TestTaskSchedulerTaskRunner : public SingleThreadTaskRunner {
  public:
   TestTaskSchedulerTaskRunner(TestTaskScheduler* task_scheduler,
                               ExecutionMode execution_mode,
@@ skipping 85 matching lines @@
 void TestTaskScheduler::JoinForTesting() {
   // TestTaskScheduler doesn't create threads so this does nothing.
 }
 
 bool TestTaskScheduler::PostTask(std::unique_ptr<internal::Task> task,
                                  const SequenceToken& sequence_token) {
   DCHECK(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);
 }
 
 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.
   task_tracker_.RunTask(std::move(task), sequence_token.IsValid()
                                              ? sequence_token
                                              : 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(
     TestTaskScheduler* task_scheduler,
     ExecutionMode execution_mode,
     TaskTraits traits)
     : task_scheduler_(task_scheduler),
       execution_mode_(execution_mode),
       sequence_token_(execution_mode == ExecutionMode::PARALLEL
                           ? SequenceToken()
@@ skipping 43 matching lines @@
   DCHECK_EQ(task_scheduler_, TaskScheduler::GetInstance());
 
   // Per contract, call JoinForTesting() before deleting the TaskScheduler.
   TaskScheduler::GetInstance()->JoinForTesting();
 
   TaskScheduler::SetInstance(nullptr);
 }
 
 }  // namespace test
 }  // namespace base