mojo: Update MessagePumpMojo to allow RunUntilIdle to purge all handles

Right now, if RunUntilIdle is called on MessageLoop that use a
MessagePumpMojo and multiple handles are ready (or multiple messages are
available on a single handle), the loop will only handle one before
returning.

This CL ensure that delegate->DoIdleWork() is only called if not handle
were ready.

Review URL: https://codereview.chromium.org/619853002

Cr-Commit-Position: refs/heads/master@{#298248}

4 files changed, 75 insertions, 11 deletions

diff --git a/mojo/common/message_pump_mojo.cc b/mojo/common/message_pump_mojo.cc
index 91b78d1..6d21af2 100644
--- a/mojo/common/message_pump_mojo.cc
+++ b/mojo/common/message_pump_mojo.cc
@@ -139,11 +139,7 @@ void MessagePumpMojo::DoRunLoop(RunState* run_state, Delegate* delegate) {
   bool more_work_is_plausible = true;
   for (;;) {
     const bool block = !more_work_is_plausible;
-    DoInternalWork(*run_state, block);
-
-    // There isn't a good way to know if there are more handles ready, we assume
-    // not.
-    more_work_is_plausible = false;
+    more_work_is_plausible = DoInternalWork(*run_state, block);
 
     if (run_state->should_quit)
       break;
@@ -166,15 +162,15 @@ void MessagePumpMojo::DoRunLoop(RunState* run_state, Delegate* delegate) {
   }
 }
 
-void MessagePumpMojo::DoInternalWork(const RunState& run_state, bool block) {
+bool MessagePumpMojo::DoInternalWork(const RunState& run_state, bool block) {
   const MojoDeadline deadline = block ? GetDeadlineForWait(run_state) : 0;
   const WaitState wait_state = GetWaitState(run_state);
   const MojoResult result =
       WaitMany(wait_state.handles, wait_state.wait_signals, deadline);
+  bool did_work = true;
   if (result == 0) {
     // Control pipe was written to.
-    uint32_t num_bytes = 0;
-    ReadMessageRaw(run_state.read_handle.get(), NULL, &num_bytes, NULL, NULL,
+    ReadMessageRaw(run_state.read_handle.get(), NULL, NULL, NULL, NULL,
                    MOJO_READ_MESSAGE_FLAG_MAY_DISCARD);
   } else if (result > 0) {
     const size_t index = static_cast<size_t>(result);
@@ -188,6 +184,7 @@ void MessagePumpMojo::DoInternalWork(const RunState& run_state, bool block) {
         RemoveFirstInvalidHandle(wait_state);
         break;
       case MOJO_RESULT_DEADLINE_EXCEEDED:
+        did_work = false;
         break;
       default:
         base::debug::Alias(&result);
@@ -208,8 +205,11 @@ void MessagePumpMojo::DoInternalWork(const RunState& run_state, bool block) {
         handlers_.find(i->first) != handlers_.end() &&
         handlers_[i->first].id == i->second.id) {
       i->second.handler->OnHandleError(i->first, MOJO_RESULT_DEADLINE_EXCEEDED);
+      handlers_.erase(i->first);
+      did_work = true;
     }
   }
+  return did_work;
 }
 
 void MessagePumpMojo::RemoveFirstInvalidHandle(const WaitState& wait_state) {
diff --git a/mojo/common/message_pump_mojo.h b/mojo/common/message_pump_mojo.h
index 3ec3457..26e28b1 100644
--- a/mojo/common/message_pump_mojo.h
+++ b/mojo/common/message_pump_mojo.h
@@ -73,8 +73,9 @@ class MOJO_COMMON_EXPORT MessagePumpMojo : public base::MessagePump {
   void DoRunLoop(RunState* run_state, Delegate* delegate);
 
   // Services the set of handles ready. If |block| is true this waits for a
-  // handle to become ready, otherwise this does not block.
-  void DoInternalWork(const RunState& run_state, bool block);
+  // handle to become ready, otherwise this does not block. Returns |true| if a
+  // handle has become ready, |false| otherwise.
+  bool DoInternalWork(const RunState& run_state, bool block);
 
   // Removes the first invalid handle. This is called if MojoWaitMany finds an
   // invalid handle.
diff --git a/mojo/common/message_pump_mojo_handler.h b/mojo/common/message_pump_mojo_handler.h
index 5809051..ba40f80 100644
--- a/mojo/common/message_pump_mojo_handler.h
+++ b/mojo/common/message_pump_mojo_handler.h
@@ -12,7 +12,7 @@ namespace mojo {
 namespace common {
 
 // Used by MessagePumpMojo to notify when a handle is either ready or has become
-// invalid.
+// invalid. In case of error, the handler will be removed.
 class MOJO_COMMON_EXPORT MessagePumpMojoHandler {
  public:
   virtual void OnHandleReady(const Handle& handle) = 0;
diff --git a/mojo/common/message_pump_mojo_unittest.cc b/mojo/common/message_pump_mojo_unittest.cc
index fb5ae24..d552942 100644
--- a/mojo/common/message_pump_mojo_unittest.cc
+++ b/mojo/common/message_pump_mojo_unittest.cc
@@ -5,6 +5,9 @@
 #include "mojo/common/message_pump_mojo.h"
 
 #include "base/message_loop/message_loop_test.h"
+#include "base/run_loop.h"
+#include "mojo/common/message_pump_mojo_handler.h"
+#include "mojo/public/cpp/system/core.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace mojo {
@@ -17,6 +20,66 @@ scoped_ptr<base::MessagePump> CreateMojoMessagePump() {
 
 RUN_MESSAGE_LOOP_TESTS(Mojo, &CreateMojoMessagePump);
 
+class CountingMojoHandler : public MessagePumpMojoHandler {
+ public:
+  CountingMojoHandler() : success_count_(0), error_count_(0) {}
+
+  virtual void OnHandleReady(const Handle& handle) override {
+    ReadMessageRaw(static_cast<const MessagePipeHandle&>(handle),
+                   NULL,
+                   NULL,
+                   NULL,
+                   NULL,
+                   MOJO_READ_MESSAGE_FLAG_NONE);
+    ++success_count_;
+  }
+  virtual void OnHandleError(const Handle& handle, MojoResult result) override {
+    ++error_count_;
+  }
+
+  int success_count() { return success_count_; }
+  int error_count() { return error_count_; }
+
+ private:
+  int success_count_;
+  int error_count_;
+
+  DISALLOW_COPY_AND_ASSIGN(CountingMojoHandler);
+};
+
+TEST(MessagePumpMojo, RunUntilIdle) {
+  base::MessageLoop message_loop(MessagePumpMojo::Create());
+  CountingMojoHandler handler;
+  MessagePipe handles;
+  MessagePumpMojo::current()->AddHandler(&handler,
+                                         handles.handle0.get(),
+                                         MOJO_HANDLE_SIGNAL_READABLE,
+                                         base::TimeTicks());
+  WriteMessageRaw(
+      handles.handle1.get(), NULL, 0, NULL, 0, MOJO_WRITE_MESSAGE_FLAG_NONE);
+  WriteMessageRaw(
+      handles.handle1.get(), NULL, 0, NULL, 0, MOJO_WRITE_MESSAGE_FLAG_NONE);
+  base::RunLoop run_loop;
+  run_loop.RunUntilIdle();
+  EXPECT_EQ(2, handler.success_count());
+}
+
+TEST(MessagePumpMojo, UnregisterAfterDeadline) {
+  base::MessageLoop message_loop(MessagePumpMojo::Create());
+  CountingMojoHandler handler;
+  MessagePipe handles;
+  MessagePumpMojo::current()->AddHandler(
+      &handler,
+      handles.handle0.get(),
+      MOJO_HANDLE_SIGNAL_READABLE,
+      base::TimeTicks::Now() - base::TimeDelta::FromSeconds(1));
+  for (int i = 0; i < 2; ++i) {
+    base::RunLoop run_loop;
+    run_loop.RunUntilIdle();
+  }
+  EXPECT_EQ(1, handler.error_count());
+}
+
 }  // namespace test
 }  // namespace common
 }  // namespace mojo