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 bca4d5f4b9733a2ca4709c4ddb8bba047ec1eb3a..e33926d6f06030c7f0d27a9cc8a663e80ee2419c 100644
--- a/chrome/test/base/ui_test_utils.cc
+++ b/chrome/test/base/ui_test_utils.cc
@@ -91,6 +91,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 = 5;
+
 namespace ui_test_utils {
 
 namespace {
@@ -104,7 +109,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,
@@ -114,12 +120,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 {
@@ -131,6 +137,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);
 };
@@ -145,7 +152,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_; }
@@ -163,7 +171,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";
         }
@@ -183,6 +191,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);
 };
@@ -261,39 +270,54 @@ 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,
+                                     const base::Closure& quit_task) {
+  MessageLoop::current()->PostTask(FROM_HERE,
+                                   MessageLoop::QuitWhenIdleClosure());
   RunMessageLoop();
-  content::BrowserThread::PostTask(thread_id, FROM_HERE,
-                                   MessageLoop::QuitClosure());
+  content::BrowserThread::PostTask(thread_id, FROM_HERE, quit_task);
 }
 
 }  // namespace
 
 void RunMessageLoop() {
-  MessageLoop* loop = MessageLoop::current();
-  MessageLoopForUI* ui_loop =
-      content::BrowserThread::CurrentlyOn(content::BrowserThread::UI) ?
-          MessageLoopForUI::current() : NULL;
-  MessageLoop::ScopedNestableTaskAllower allow(loop);
-  if (ui_loop) {
-#if defined(USE_AURA)
-    ui_loop->Run();
-#elif defined(TOOLKIT_VIEWS)
+  base::RunLoop run_loop;
+  RunRunLoop(run_loop.AsWeakPtr());
+}
+
+void RunRunLoop(const base::WeakPtr<base::RunLoop>& run_loop) {

[jar (doing other things), 2012/06/23 02:43:19]

nit: I didn't like the RunRun... prefix.  Perhaps better would be:
RunThisRunLoop
or
RunWeakRunLoop
or
RunWeakPtrRunLoop
or
other??

In truth, I also really wished this was a method on this weak-ptr... and then it
would just be a Run() method.  Is it worth making a local helper class
containing the weak pointer? Is that too clumsy? 

[jbates, 2012/06/25 20:04:57]

Renamed to RunThisRunLoop. I tried the helper class, but it didn't seem to be
worth its code-weight.

+  CHECK(run_loop);
+  MessageLoop::ScopedNestableTaskAllower allow(MessageLoop::current());
+  if (content::BrowserThread::CurrentlyOn(content::BrowserThread::UI)) {
+#if !defined(USE_AURA) && defined(TOOLKIT_VIEWS)
     views::AcceleratorHandler handler;
-    ui_loop->RunWithDispatcher(&handler);
-#elif defined(OS_POSIX) && !defined(OS_MACOSX) && !defined(OS_ANDROID)
-    ui_loop->RunWithDispatcher(NULL);
-#else
-    ui_loop->Run();
+    run_loop->set_dispatcher(&handler);
 #endif
+    run_loop->Run();
+  } else {
+    run_loop->Run();

[jar (doing other things), 2012/06/23 02:43:19]

Line 296 and 298 can now be pulled out of the if blocks, and then in fact, line
291 can now go inside the #if.

[jbates, 2012/06/25 20:04:57]

Done.

+  }
+}
+
+static void DeferredQuitRunLoop(
+    const base::WeakPtr<base::RunLoop>& run_loop,
+    int num_quit_deferrals) {
+  if (num_quit_deferrals <= 0) {
+    if (run_loop) run_loop->Quit();

[jar (doing other things), 2012/06/23 02:43:19]

Do we gain value by iterating and posting when we don't have a run_loop any
more, or can we just return once it is gone?  After all, when we're done posting
(via line 308) and running, we'll just give up and do nothing anyway.
Could we just:
if (!run_loop)
  return;
before line 305?

It seems that if the run loop is gone, we'll just be posting into a messageloop
that has exited, and building cruft tasks there.

Or can we assert that we are the only way for this loop to finally end, and the
run_loop is always still intact?

[jbates, 2012/06/25 20:04:57]

Done.

   } else {
-    loop->Run();
+    MessageLoop::current()->PostTask(FROM_HERE,
+        base::Bind(&DeferredQuitRunLoop, run_loop, num_quit_deferrals - 1));
   }
 }
 
+// Quit the MessageLoop Run of the specified |run_id|.

[jar (doing other things), 2012/06/23 02:43:19]

nit: update comment "run_id"

[jbates, 2012/06/25 20:04:57]

Done.

+void QuitRunLoop(const base::WeakPtr<base::RunLoop>& run_loop) {

[jar (doing other things), 2012/06/23 02:43:19]

This name is not as jarring... but if you did encapsulate the WeakPtr to a
RunLoop, this would indeed just be a Quit method... or perhaps a DeferredQuit()
method.

[jbates, 2012/06/25 20:04:57]

Changed to QuitThisRunLoop to be consistent with RunThisRunLoop.

+  DeferredQuitRunLoop(run_loop, kNumQuitDeferrals);
+}
+
 void RunAllPendingInMessageLoop() {
-  MessageLoop::current()->PostTask(FROM_HERE, MessageLoop::QuitClosure());
+  MessageLoop::current()->PostTask(FROM_HERE,
+                                   MessageLoop::QuitWhenIdleClosure());
   ui_test_utils::RunMessageLoop();
 }
 
@@ -307,10 +331,12 @@ void RunAllPendingInMessageLoop(content::BrowserThread::ID thread_id) {
     NOTREACHED();
     return;
   }
-  content::BrowserThread::PostTask(thread_id, FROM_HERE,
-      base::Bind(&RunAllPendingMessageAndSendQuit, current_thread_id));
 
-  ui_test_utils::RunMessageLoop();
+  base::RunLoop run_loop;
+  content::BrowserThread::PostTask(thread_id, FROM_HERE,
+      base::Bind(&RunAllPendingMessageAndSendQuit, current_thread_id,
+                 run_loop.QuitClosure()));
+  ui_test_utils::RunRunLoop(run_loop.AsWeakPtr());
 }
 
 bool GetCurrentTabTitle(const Browser* browser, string16* title) {
@@ -329,10 +355,10 @@ void WaitForNavigations(NavigationController* controller,
   content::TestNavigationObserver observer(
       content::Source<NavigationController>(controller), NULL,
       number_of_navigations);
+  base::RunLoop run_loop;
   observer.WaitForObservation(
-      base::Bind(&ui_test_utils::RunMessageLoop),
-      base::Bind(&MessageLoop::Quit,
-                 base::Unretained(MessageLoopForUI::current())));
+      base::Bind(&ui_test_utils::RunRunLoop, run_loop.AsWeakPtr()),
+      base::Bind(&ui_test_utils::QuitRunLoop, run_loop.AsWeakPtr()));
 }
 
 void WaitForNewTab(Browser* browser) {
@@ -382,11 +408,10 @@ void NavigateToURL(browser::NavigateParams* params) {
   content::TestNavigationObserver observer(
       content::NotificationService::AllSources(), NULL, 1);
   browser::Navigate(params);
+  base::RunLoop run_loop;
   observer.WaitForObservation(
-      base::Bind(&ui_test_utils::RunMessageLoop),
-      base::Bind(&MessageLoop::Quit,
-                 base::Unretained(MessageLoopForUI::current())));
-
+      base::Bind(&ui_test_utils::RunRunLoop, run_loop.AsWeakPtr()),
+      base::Bind(&ui_test_utils::QuitRunLoop, run_loop.AsWeakPtr()));
 }
 
 void NavigateToURL(Browser* browser, const GURL& url) {
@@ -456,10 +481,10 @@ static void NavigateToURLWithDispositionBlockUntilNavigationsComplete(
     web_contents = browser->GetActiveWebContents();
   }
   if (disposition == CURRENT_TAB) {
+    base::RunLoop run_loop;
     same_tab_observer.WaitForObservation(
-        base::Bind(&ui_test_utils::RunMessageLoop),
-        base::Bind(&MessageLoop::Quit,
-                   base::Unretained(MessageLoopForUI::current())));
+        base::Bind(&ui_test_utils::RunRunLoop, run_loop.AsWeakPtr()),
+        base::Bind(&ui_test_utils::QuitRunLoop, run_loop.AsWeakPtr()));
     return;
   } else if (web_contents) {
     NavigationController* controller = &web_contents->GetController();
@@ -703,9 +728,10 @@ void RegisterAndWait(content::NotificationObserver* observer,
 void WaitForBookmarkModelToLoad(BookmarkModel* model) {
   if (model->IsLoaded())
     return;
-  BookmarkLoadObserver observer;
+  base::RunLoop run_loop;
+  BookmarkLoadObserver observer(run_loop.AsWeakPtr());
   model->AddObserver(&observer);
-  RunMessageLoop();
+  RunRunLoop(run_loop.AsWeakPtr());
   model->RemoveObserver(&observer);
   ASSERT_TRUE(model->IsLoaded());
 }
@@ -763,14 +789,15 @@ bool SendKeyPressSync(const Browser* browser,
   gfx::NativeWindow window = NULL;
   if (!GetNativeWindow(browser, &window))
     return false;
-
+  scoped_refptr<MessageLoopRunner> runner =
+      new MessageLoopRunner;

[jar (doing other things), 2012/06/23 02:43:19]

nit: no need to wrap.

[jbates, 2012/06/25 20:04:57]

Done.

   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) {
@@ -781,7 +808,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();
 }
 
@@ -803,51 +830,43 @@ 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;

[jar (doing other things), 2012/06/23 02:43:19]

nit: un-wrap (here and below in file, like line 845, 1276, etc.)

[jbates, 2012/06/25 20:04:57]

Done.

+  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() {
 }
 
-TimedMessageLoopRunner::~TimedMessageLoopRunner() {
-  if (owned_)
-    delete loop_;
+MessageLoopRunner::~MessageLoopRunner() {
 }
 
-void TimedMessageLoopRunner::RunFor(int ms) {
-  QuitAfter(ms);
-  quit_loop_invoked_ = false;
-  loop_->Run();
+void MessageLoopRunner::Run() {
+  ui_test_utils::RunRunLoop(run_loop_.AsWeakPtr());
 }
 
-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() {
+  ui_test_utils::QuitRunLoop(run_loop_.AsWeakPtr());
 }
 
 TestWebSocketServer::TestWebSocketServer()
@@ -1012,7 +1031,8 @@ void WindowedNotificationObserver::Wait() {
     return;
 
   running_ = true;
-  ui_test_utils::RunMessageLoop();
+  message_loop_runner_ = new MessageLoopRunner;
+  message_loop_runner_->Run();
   EXPECT_TRUE(seen_);
 }
 
@@ -1026,7 +1046,7 @@ void WindowedNotificationObserver::Observe(
   if (!running_)
     return;
 
-  MessageLoopForUI::current()->Quit();
+  message_loop_runner_->QuitNow();
   running_ = false;
 }
 
@@ -1077,7 +1097,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_;
 }
 
@@ -1103,8 +1124,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()
@@ -1139,7 +1163,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();
   }
 }
 
@@ -1151,7 +1175,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())
@@ -1179,7 +1204,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_;
   }
 
@@ -1214,13 +1240,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);
 };
@@ -1246,9 +1273,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 {