Add support to base::Timer for custom task runners.

base/timer/timer_unittest.cc

Index: base/timer/timer_unittest.cc
diff --git a/base/timer/timer_unittest.cc b/base/timer/timer_unittest.cc
index 0fb2b454ad764f7497dca7878910bee3ca701238..c56bc1dce59a1b53def65d6f780f29d28a76f468 100644
--- a/base/timer/timer_unittest.cc
+++ b/base/timer/timer_unittest.cc
@@ -4,10 +4,12 @@
 
 #include "base/memory/scoped_ptr.h"
 #include "base/message_loop/message_loop.h"
+#include "base/test/test_simple_task_runner.h"
 #include "base/timer/timer.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 using base::TimeDelta;
+using base::SingleThreadTaskRunner;
 
 namespace {
 
@@ -26,20 +28,28 @@ class OneShotTimerTester {
  public:
   explicit OneShotTimerTester(bool* did_run, unsigned milliseconds = 10)
       : did_run_(did_run),
-        delay_ms_(milliseconds) {
+        delay_ms_(milliseconds),
+        quit_message_loop_(true) {
   }
   void Start() {
     timer_.Start(FROM_HERE, TimeDelta::FromMilliseconds(delay_ms_), this,
                  &OneShotTimerTester::Run);
   }
+  void SetTaskRunner(scoped_refptr<SingleThreadTaskRunner> task_runner) {
+    quit_message_loop_ = false;
+    timer_.SetTaskRunner(task_runner);
+  }
  private:
   void Run() {
     *did_run_ = true;
-    base::MessageLoop::current()->QuitWhenIdle();
+    if (quit_message_loop_) {
+      base::MessageLoop::current()->QuitWhenIdle();
+    }
   }
   bool* did_run_;
   base::OneShotTimer<OneShotTimerTester> timer_;
   const unsigned delay_ms_;
+  bool quit_message_loop_;
 };
 
 class OneShotSelfDeletingTimerTester {
@@ -132,6 +142,52 @@ void RunTest_OneShotSelfDeletingTimer(
   EXPECT_TRUE(did_run);
 }
 
+void RunTest_OneShotTimer_CustomTaskRunner(
+    base::MessageLoop::Type message_loop_type) {
+  base::MessageLoop loop(message_loop_type);
+  scoped_refptr<base::TestSimpleTaskRunner> task_runner =
+      new base::TestSimpleTaskRunner();
+
+  // Create three timers. While the second one uses the message loop (default
+  // behavior), the first and third one use the task runner.
+  bool did_run1 = false;
+  bool did_run2 = false;
+  bool did_run3 = false;
+  OneShotTimerTester f1(&did_run1);
+  OneShotTimerTester f2(&did_run2);
+  OneShotTimerTester f3(&did_run3);
+  f1.SetTaskRunner(task_runner);
+  f3.SetTaskRunner(task_runner);
+
+  // Start the first and second timer and run the task runner. Only the first
+  // timer should go off.
+  f1.Start();
+  f2.Start();
+  task_runner->RunUntilIdle();
+
+  EXPECT_TRUE(did_run1);
+  EXPECT_FALSE(did_run2);
+  EXPECT_FALSE(did_run3);
+  did_run1 = false;
+
+  // Start the third timer and run the message loop. The second timer should
+  // go off this time.
+  f3.Start();
+  base::MessageLoop::current()->Run();
+
+  EXPECT_FALSE(did_run1);
+  EXPECT_TRUE(did_run2);
+  EXPECT_FALSE(did_run3);
+  did_run2 = false;
+
+  // Run the task runner again. Finally, the third timer should go off.
+  task_runner->RunUntilIdle();
+
+  EXPECT_FALSE(did_run1);
+  EXPECT_FALSE(did_run2);
+  EXPECT_TRUE(did_run3);
+}
+
 void RunTest_RepeatingTimer(base::MessageLoop::Type message_loop_type,
                             const TimeDelta& delay) {
   base::MessageLoop loop(message_loop_type);
@@ -310,6 +366,12 @@ TEST(TimerTest, OneShotSelfDeletingTimer) {
   }
 }
 
+TEST(TimerTest, OneShotTimer_CustomTaskRunner) {
+  for (int i = 0; i < kNumTestingMessageLoops; i++) {
+    RunTest_OneShotTimer_CustomTaskRunner(testing_message_loops[i]);
+  }
+}
+
 TEST(TimerTest, RepeatingTimer) {
   for (int i = 0; i < kNumTestingMessageLoops; i++) {
     RunTest_RepeatingTimer(testing_message_loops[i],