Add base::RunLoop and update ui_test_utils to use it to reduce flakiness

chrome/test/base/ui_test_utils.cc

Index: chrome/test/base/ui_test_utils.cc
diff --git a/chrome/test/base/ui_test_utils.cc b/chrome/test/base/ui_test_utils.cc
index f1b27b0b8c56af5ffc3f7f6c6b5762b78f186889..bf4ab0df99329bb3287ab542aea1ee16312018a2 100644
--- a/chrome/test/base/ui_test_utils.cc
+++ b/chrome/test/base/ui_test_utils.cc
@@ -92,6 +92,11 @@ using content::WebContents;
 
 static const int kDefaultWsPort = 8880;
 
+// Number of times to repost a QuitNow task so that the MessageLoop finishes up
+// pending tasks and tasks posted by those pending tasks without risking the
+// potential hang behavior of MessageLoop::QuitWhenIdle.
+static const int kNumQuitDeferrals = 20;
+
 namespace ui_test_utils {
 
 namespace {
@@ -105,7 +110,8 @@ class DOMOperationObserver : public content::NotificationObserver,
         did_respond_(false) {
     registrar_.Add(this, content::NOTIFICATION_DOM_OPERATION_RESPONSE,
                    content::Source<RenderViewHost>(render_view_host));
-    ui_test_utils::RunMessageLoop();
+    message_loop_runner_ = new MessageLoopRunner;
+    message_loop_runner_->Run();
   }
 
   virtual void Observe(int type,
@@ -115,12 +121,12 @@ class DOMOperationObserver : public content::NotificationObserver,
     content::Details<DomOperationNotificationDetails> dom_op_details(details);
     response_ = dom_op_details->json;
     did_respond_ = true;
-    MessageLoopForUI::current()->Quit();
+    message_loop_runner_->QuitNow();
   }
 
   // Overridden from content::WebContentsObserver:
   virtual void RenderViewGone(base::TerminationStatus status) OVERRIDE {
-    MessageLoopForUI::current()->Quit();
+    message_loop_runner_->QuitNow();
   }
 
   bool GetResponse(std::string* response) WARN_UNUSED_RESULT {
@@ -132,6 +138,7 @@ class DOMOperationObserver : public content::NotificationObserver,
   content::NotificationRegistrar registrar_;
   std::string response_;
   bool did_respond_;
+  scoped_refptr<MessageLoopRunner> message_loop_runner_;
 
   DISALLOW_COPY_AND_ASSIGN(DOMOperationObserver);
 };
@@ -146,7 +153,8 @@ class FindInPageNotificationObserver : public content::NotificationObserver {
         parent_tab->find_tab_helper()->current_find_request_id();
     registrar_.Add(this, chrome::NOTIFICATION_FIND_RESULT_AVAILABLE,
                    content::Source<WebContents>(parent_tab_->web_contents()));
-    ui_test_utils::RunMessageLoop();
+    message_loop_runner_ = new MessageLoopRunner;
+    message_loop_runner_->Run();
   }
 
   int active_match_ordinal() const { return active_match_ordinal_; }
@@ -164,7 +172,7 @@ class FindInPageNotificationObserver : public content::NotificationObserver {
           active_match_ordinal_ = find_details->active_match_ordinal();
         if (find_details->final_update()) {
           number_of_matches_ = find_details->number_of_matches();
-          MessageLoopForUI::current()->Quit();
+          message_loop_runner_->QuitNow();
         } else {
           DVLOG(1) << "Ignoring, since we only care about the final message";
         }
@@ -184,6 +192,7 @@ class FindInPageNotificationObserver : public content::NotificationObserver {
   // The id of the current find request, obtained from WebContents. Allows us
   // to monitor when the search completes.
   int current_find_request_id_;
+  scoped_refptr<MessageLoopRunner> message_loop_runner_;
 
   DISALLOW_COPY_AND_ASSIGN(FindInPageNotificationObserver);
 };
@@ -262,16 +271,22 @@ bool ExecuteJavaScriptHelper(RenderViewHost* render_view_host,
   return true;
 }
 
-void RunAllPendingMessageAndSendQuit(content::BrowserThread::ID thread_id) {
-  MessageLoop::current()->PostTask(FROM_HERE, MessageLoop::QuitClosure());
+void RunAllPendingMessageAndSendQuit(content::BrowserThread::ID thread_id,
+                                     MessageLoop::RunID message_loop_run_id) {
+  MessageLoop::current()->PostTask(FROM_HERE,
+                                   MessageLoop::QuitWhenIdleClosure());
   RunMessageLoop();
   content::BrowserThread::PostTask(thread_id, FROM_HERE,
-                                   MessageLoop::QuitClosure());
+      MessageLoop::QuitNowWithIDClosure(message_loop_run_id));
 }
 
 }  // namespace
 
 void RunMessageLoop() {
+  RunMessageLoopWithID(MessageLoop::current()->GetRunID());
+}
+
+void RunMessageLoopWithID(const MessageLoop::RunID& run_id) {
   MessageLoop* loop = MessageLoop::current();
   MessageLoopForUI* ui_loop =
       content::BrowserThread::CurrentlyOn(content::BrowserThread::UI) ?
@@ -279,22 +294,45 @@ void RunMessageLoop() {
   MessageLoop::ScopedNestableTaskAllower allow(loop);
   if (ui_loop) {
 #if defined(USE_AURA)
-    ui_loop->Run();
+    ui_loop->RunWithID(run_id);
 #elif defined(TOOLKIT_VIEWS)
     views::AcceleratorHandler handler;
-    ui_loop->RunWithDispatcher(&handler);
+    ui_loop->RunWithDispatcher(&handler, run_id);
 #elif defined(OS_POSIX) && !defined(OS_MACOSX) && !defined(OS_ANDROID)
-    ui_loop->RunWithDispatcher(NULL);
+    ui_loop->RunWithDispatcher(NULL, run_id);
 #else
-    ui_loop->Run();
+    ui_loop->RunWithID(run_id);
 #endif
   } else {
-    loop->Run();
+    loop->RunWithID(run_id);
+  }
+}
+
+// Quit the MessageLoop Run of the specified |run_id|.
+void QuitMessageLoopWithID(const MessageLoop::RunID& run_id) {
+  MessageLoop::current()->PostTask(
+      FROM_HERE, ui_test_utils::QuitMessageLoopWithIDClosure(run_id));
+}
+
+static void DeferTask(const base::Closure& task,
+                      int num_post_task_deferrals) {
+  if (num_post_task_deferrals <= 0) {
+    MessageLoop::current()->PostTask(FROM_HERE, task);
+  } else {
+    MessageLoop::current()->PostTask(
+        FROM_HERE, base::Bind(&DeferTask, task, num_post_task_deferrals - 1));
   }
 }
 
+base::Closure QuitMessageLoopWithIDClosure(const MessageLoop::RunID& run_id) {
+  return base::Bind(&DeferTask,
+                    MessageLoop::QuitNowWithIDClosure(run_id),
+                    kNumQuitDeferrals);
+}
+
 void RunAllPendingInMessageLoop() {
-  MessageLoop::current()->PostTask(FROM_HERE, MessageLoop::QuitClosure());
+  MessageLoop::current()->PostTask(FROM_HERE,
+                                   MessageLoop::QuitWhenIdleClosure());
   ui_test_utils::RunMessageLoop();
 }
 
@@ -308,10 +346,12 @@ void RunAllPendingInMessageLoop(content::BrowserThread::ID thread_id) {
     NOTREACHED();
     return;
   }
+
+  MessageLoop::RunID run_id = MessageLoop::current()->GetRunID();
   content::BrowserThread::PostTask(thread_id, FROM_HERE,
-      base::Bind(&RunAllPendingMessageAndSendQuit, current_thread_id));
+      base::Bind(&RunAllPendingMessageAndSendQuit, current_thread_id, run_id));
 
-  ui_test_utils::RunMessageLoop();
+  ui_test_utils::RunMessageLoopWithID(run_id);
 }
 
 bool GetCurrentTabTitle(const Browser* browser, string16* title) {
@@ -330,10 +370,10 @@ void WaitForNavigations(NavigationController* controller,
   content::TestNavigationObserver observer(
       content::Source<NavigationController>(controller), NULL,
       number_of_navigations);
+  MessageLoop::RunID message_loop_run_id = MessageLoop::current()->GetRunID();
   observer.WaitForObservation(
-      base::Bind(&ui_test_utils::RunMessageLoop),
-      base::Bind(&MessageLoop::Quit,
-                 base::Unretained(MessageLoopForUI::current())));
+      base::Bind(&ui_test_utils::RunMessageLoopWithID, message_loop_run_id),
+      base::Bind(&ui_test_utils::QuitMessageLoopWithID, message_loop_run_id));
 }
 
 void WaitForNewTab(Browser* browser) {
@@ -397,11 +437,10 @@ void NavigateToURL(browser::NavigateParams* params) {
   content::TestNavigationObserver observer(
       content::NotificationService::AllSources(), NULL, 1);
   browser::Navigate(params);
+  MessageLoop::RunID message_loop_run_id = MessageLoop::current()->GetRunID();
   observer.WaitForObservation(
-      base::Bind(&ui_test_utils::RunMessageLoop),
-      base::Bind(&MessageLoop::Quit,
-                 base::Unretained(MessageLoopForUI::current())));
-
+      base::Bind(&ui_test_utils::RunMessageLoopWithID, message_loop_run_id),
+      base::Bind(&ui_test_utils::QuitMessageLoopWithID, message_loop_run_id));
 }
 
 void NavigateToURL(Browser* browser, const GURL& url) {
@@ -471,10 +510,10 @@ static void NavigateToURLWithDispositionBlockUntilNavigationsComplete(
     web_contents = browser->GetActiveWebContents();
   }
   if (disposition == CURRENT_TAB) {
+    MessageLoop::RunID message_loop_run_id = MessageLoop::current()->GetRunID();
     same_tab_observer.WaitForObservation(
-        base::Bind(&ui_test_utils::RunMessageLoop),
-        base::Bind(&MessageLoop::Quit,
-                   base::Unretained(MessageLoopForUI::current())));
+        base::Bind(&ui_test_utils::RunMessageLoopWithID, message_loop_run_id),
+        base::Bind(&ui_test_utils::QuitMessageLoopWithID, message_loop_run_id));
     return;
   } else if (web_contents) {
     NavigationController* controller = &web_contents->GetController();
@@ -718,9 +757,10 @@ void RegisterAndWait(content::NotificationObserver* observer,
 void WaitForBookmarkModelToLoad(BookmarkModel* model) {
   if (model->IsLoaded())
     return;
-  BookmarkLoadObserver observer;
+  MessageLoop::RunID run_id = MessageLoop::current()->GetRunID();
+  BookmarkLoadObserver observer(run_id);
   model->AddObserver(&observer);
-  RunMessageLoop();
+  RunMessageLoopWithID(run_id);
   model->RemoveObserver(&observer);
   ASSERT_TRUE(model->IsLoaded());
 }
@@ -778,14 +818,15 @@ bool SendKeyPressSync(const Browser* browser,
   gfx::NativeWindow window = NULL;
   if (!GetNativeWindow(browser, &window))
     return false;
-
+  scoped_refptr<MessageLoopRunner> runner =
+      new MessageLoopRunner;
   bool result;
   result = ui_controls::SendKeyPressNotifyWhenDone(
-      window, key, control, shift, alt, command, MessageLoop::QuitClosure());
+      window, key, control, shift, alt, command, runner->QuitNowClosure());
 #if defined(OS_WIN)
   if (!result && BringBrowserWindowToFront(browser)) {
     result = ui_controls::SendKeyPressNotifyWhenDone(
-        window, key, control, shift, alt, command, MessageLoop::QuitClosure());
+        window, key, control, shift, alt, command, runner->QuitNowClosure());
   }
 #endif
   if (!result) {
@@ -796,7 +837,7 @@ bool SendKeyPressSync(const Browser* browser,
   // Run the message loop. It'll stop running when either the key was received
   // or the test timed out (in which case testing::Test::HasFatalFailure should
   // be set).
-  RunMessageLoop();
+  runner->Run();
   return !testing::Test::HasFatalFailure();
 }
 
@@ -818,51 +859,57 @@ bool SendKeyPressAndWait(const Browser* browser,
 }
 
 bool SendMouseMoveSync(const gfx::Point& location) {
-  if (!ui_controls::SendMouseMoveNotifyWhenDone(location.x(), location.y(),
-                                                MessageLoop::QuitClosure())) {
+  scoped_refptr<MessageLoopRunner> runner =
+      new MessageLoopRunner;
+  if (!ui_controls::SendMouseMoveNotifyWhenDone(
+          location.x(), location.y(), runner->QuitNowClosure())) {
     return false;
   }
-  RunMessageLoop();
+  runner->Run();
   return !testing::Test::HasFatalFailure();
 }
 
 bool SendMouseEventsSync(ui_controls::MouseButton type, int state) {
+  scoped_refptr<MessageLoopRunner> runner =
+      new MessageLoopRunner;
   if (!ui_controls::SendMouseEventsNotifyWhenDone(
-          type, state, MessageLoop::QuitClosure())) {
+          type, state, runner->QuitNowClosure())) {
     return false;
   }
-  RunMessageLoop();
+  runner->Run();
   return !testing::Test::HasFatalFailure();
 }
 
-TimedMessageLoopRunner::TimedMessageLoopRunner()
-    : loop_(new MessageLoopForUI()),
-      owned_(true),
-      quit_loop_invoked_(false) {
+MessageLoopRunner::MessageLoopRunner()
+    : message_loop_run_id_(MessageLoop::current()->GetRunID()),
+      run_called_(false),
+      quit_received_(false),
+      running_(false) {
 }
 
-TimedMessageLoopRunner::~TimedMessageLoopRunner() {
-  if (owned_)
-    delete loop_;
+MessageLoopRunner::~MessageLoopRunner() {
 }
 
-void TimedMessageLoopRunner::RunFor(int ms) {
-  QuitAfter(ms);
-  quit_loop_invoked_ = false;
-  loop_->Run();
+void MessageLoopRunner::Run() {
+  CHECK(!run_called_);
+  run_called_ = true;
+
+  if (quit_received_)
+    return;
+
+  running_ = true;
+  ui_test_utils::RunMessageLoopWithID(message_loop_run_id_);
+  running_ = false;
 }
 
-void TimedMessageLoopRunner::Quit() {
-  quit_loop_invoked_ = true;
-  loop_->PostTask(FROM_HERE, MessageLoop::QuitClosure());
+base::Closure MessageLoopRunner::QuitNowClosure() {
+  return base::Bind(&MessageLoopRunner::QuitNow, this);
 }
 
-void TimedMessageLoopRunner::QuitAfter(int ms) {
-  quit_loop_invoked_ = true;
-  loop_->PostDelayedTask(
-      FROM_HERE,
-      MessageLoop::QuitClosure(),
-      base::TimeDelta::FromMilliseconds(ms));
+void MessageLoopRunner::QuitNow() {
+  quit_received_ = true;
+  if (running_)
+    ui_test_utils::QuitMessageLoopWithID(message_loop_run_id_);
 }
 
 TestWebSocketServer::TestWebSocketServer()
@@ -1027,7 +1074,8 @@ void WindowedNotificationObserver::Wait() {
     return;
 
   running_ = true;
-  ui_test_utils::RunMessageLoop();
+  message_loop_runner_ = new MessageLoopRunner;
+  message_loop_runner_->Run();
   EXPECT_TRUE(seen_);
 }
 
@@ -1041,7 +1089,7 @@ void WindowedNotificationObserver::Observe(
   if (!running_)
     return;
 
-  MessageLoopForUI::current()->Quit();
+  message_loop_runner_->QuitNow();
   running_ = false;
 }
 
@@ -1092,7 +1140,8 @@ const string16& TitleWatcher::WaitAndGetTitle() {
   if (expected_title_observed_)
     return observed_title_;
   quit_loop_on_observation_ = true;
-  ui_test_utils::RunMessageLoop();
+  message_loop_runner_ = new MessageLoopRunner;
+  message_loop_runner_->Run();
   return observed_title_;
 }
 
@@ -1118,8 +1167,11 @@ void TitleWatcher::Observe(int type,
     return;
   observed_title_ = *it;
   expected_title_observed_ = true;
-  if (quit_loop_on_observation_)
-    MessageLoopForUI::current()->Quit();
+  if (quit_loop_on_observation_) {
+    // Only call Quit once, on first Observe:
+    quit_loop_on_observation_ = false;
+    message_loop_runner_->QuitNow();
+  }
 }
 
 BrowserAddedObserver::BrowserAddedObserver()
@@ -1154,7 +1206,7 @@ void DOMMessageQueue::Observe(int type,
   message_queue_.push(dom_op_details->json);
   if (waiting_for_message_) {
     waiting_for_message_ = false;
-    MessageLoopForUI::current()->Quit();
+    message_loop_runner_->QuitNow();
   }
 }
 
@@ -1166,7 +1218,8 @@ bool DOMMessageQueue::WaitForMessage(std::string* message) {
   if (message_queue_.empty()) {
     waiting_for_message_ = true;
     // This will be quit when a new message comes in.
-    RunMessageLoop();
+    message_loop_runner_ = new MessageLoopRunner;
+    message_loop_runner_->Run();
   }
   // The queue should not be empty, unless we were quit because of a timeout.
   if (message_queue_.empty())
@@ -1194,7 +1247,8 @@ class SnapshotTaker {
         base::Bind(&SnapshotTaker::OnSnapshotTaken, base::Unretained(this)),
         page_size,
         desired_size);
-    ui_test_utils::RunMessageLoop();
+    message_loop_runner_ = new MessageLoopRunner;
+    message_loop_runner_->Run();
     return snapshot_taken_;
   }
 
@@ -1229,13 +1283,14 @@ class SnapshotTaker {
   void OnSnapshotTaken(const SkBitmap& bitmap) {
     *bitmap_ = bitmap;
     snapshot_taken_ = true;
-    MessageLoop::current()->Quit();
+    message_loop_runner_->QuitNow();
   }
 
   SkBitmap* bitmap_;
   // Whether the snapshot was actually taken and received by this SnapshotTaker.
   // This will be false if the test times out.
   bool snapshot_taken_;
+  scoped_refptr<MessageLoopRunner> message_loop_runner_;
 
   DISALLOW_COPY_AND_ASSIGN(SnapshotTaker);
 };
@@ -1261,9 +1316,11 @@ void OverrideGeolocation(double latitude, double longitude) {
   position.altitude = 0.;
   position.accuracy = 0.;
   position.timestamp = base::Time::Now();
+  scoped_refptr<MessageLoopRunner> runner =
+      new MessageLoopRunner;
   content::OverrideLocationForTesting(position,
-                                      base::Bind(MessageLoop::QuitClosure()));
-  RunMessageLoop();
+                                      runner->QuitNowClosure());
+  runner->Run();
 }
 
 namespace internal {