Implement MessagePumpMojo using WaitSet.

mojo/message_pump/message_pump_mojo.cc

 // Copyright 2013 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 "mojo/message_pump/message_pump_mojo.h"
 
 #include <algorithm>
+#include <map>
 #include <vector>
 
 #include "base/debug/alias.h"
 #include "base/lazy_instance.h"
 #include "base/logging.h"
 #include "base/threading/thread_local.h"
 #include "base/time/time.h"
 #include "mojo/message_pump/message_pump_mojo_handler.h"
 #include "mojo/message_pump/time_helper.h"
+#include "mojo/public/c/system/wait_set.h"
 
 namespace mojo {
 namespace common {
 namespace {
 
 base::LazyInstance<base::ThreadLocalPointer<MessagePumpMojo> >::Leaky
     g_tls_current_pump = LAZY_INSTANCE_INITIALIZER;
 
 MojoDeadline TimeTicksToMojoDeadline(base::TimeTicks time_ticks,
                                      base::TimeTicks now) {
   // The is_null() check matches that of HandleWatcher as well as how
   // |delayed_work_time| is used.
   if (time_ticks.is_null())
     return MOJO_DEADLINE_INDEFINITE;
   const int64_t delta = (time_ticks - now).InMicroseconds();
   return delta < 0 ? static_cast<MojoDeadline>(0) :
                      static_cast<MojoDeadline>(delta);
 }
 
 }  // namespace
 
-// State needed for one iteration of WaitMany. The first handle and flags
-// corresponds to that of the control pipe.
-struct MessagePumpMojo::WaitState {
-  std::vector<Handle> handles;
-  std::vector<MojoHandleSignals> wait_signals;
-};
-
 struct MessagePumpMojo::RunState {
   RunState() : should_quit(false) {}
 
   base::TimeTicks delayed_work_time;
 
   bool should_quit;
 };
 
 MessagePumpMojo::MessagePumpMojo() : run_state_(NULL), next_handler_id_(0) {
   DCHECK(!current())
       << "There is already a MessagePumpMojo instance on this thread.";
   g_tls_current_pump.Pointer()->Set(this);
 
   MojoResult result = CreateMessagePipe(nullptr, &read_handle_, &write_handle_);
   CHECK_EQ(result, MOJO_RESULT_OK);
   CHECK(read_handle_.is_valid());
   CHECK(write_handle_.is_valid());
+
+  MojoHandle handle;
+  result = MojoCreateWaitSet(&handle);
+  CHECK_EQ(result, MOJO_RESULT_OK);
+  wait_set_handle_.reset(Handle(handle));
+  CHECK(wait_set_handle_.is_valid());
+
+  result = MojoAddHandle(wait_set_handle_.get().value(),
+                         read_handle_.get().value(),
+                         MOJO_HANDLE_SIGNAL_READABLE);
+  CHECK_EQ(result, MOJO_RESULT_OK);
 }
 
 MessagePumpMojo::~MessagePumpMojo() {
   DCHECK_EQ(this, current());
   g_tls_current_pump.Pointer()->Set(NULL);
 }
 
 // static
 scoped_ptr<base::MessagePump> MessagePumpMojo::Create() {
   return scoped_ptr<MessagePump>(new MessagePumpMojo());
@@ skipping 15 matching lines @@
   Handler handler_data;
   handler_data.handler = handler;
   handler_data.wait_signals = wait_signals;
   handler_data.deadline = deadline;
   handler_data.id = next_handler_id_++;
   handlers_[handle] = handler_data;
   if (!deadline.is_null()) {
     bool inserted = deadline_handles_.insert(handle).second;
     DCHECK(inserted);
   }
+
+  MojoResult result = MojoAddHandle(wait_set_handle_.get().value(),
+                                    handle.value(), wait_signals);
+  // Because stopping a HandleWatcher is now asynchronous, it's possible for the
+  // handle to no longer be open at this point.
+  CHECK(result == MOJO_RESULT_OK || result == MOJO_RESULT_INVALID_ARGUMENT);

[sky, 2015/12/11 17:49:01]

Can you add test coverage of the add with closed pipe? I don't think we have
that as previously we did not allow adding closed handles.

[Anand Mistry (off Chromium), 2015/12/14 03:16:20]

Hm. Here's a question for you. If you add a closed handle, should it result in
the error handler being called in the next loop iteration, or just be ignored
and assume the caller will remove it. Currently, the latter happens. But maybe
this function should be more proactive.

[sky, 2015/12/28 17:29:44]

Is there a reason you don't want to notify immediately?

 }
 
 void MessagePumpMojo::RemoveHandler(const Handle& handle) {
+  MojoResult result = MojoRemoveHandle(wait_set_handle_.get().value(),
+                                       handle.value());
+  // At this point, it's possible that the handle has been closed, which would
+  // cause MojoRemoveHandle() to return MOJO_RESULT_INVALID_ARGUMENT. It's also
+  // possible for the handle to have already been removed, so all of the
+  // possible error codes are valid here.
+  CHECK(result == MOJO_RESULT_OK || result == MOJO_RESULT_NOT_FOUND ||
+        result == MOJO_RESULT_INVALID_ARGUMENT);
+
   handlers_.erase(handle);
   deadline_handles_.erase(handle);
 }
 
 void MessagePumpMojo::AddObserver(Observer* observer) {
   observers_.AddObserver(observer);
 }
 
 void MessagePumpMojo::RemoveObserver(Observer* observer) {
   observers_.RemoveObserver(observer);
@@ skipping 53 matching lines @@
     if (more_work_is_plausible)
       continue;
 
     more_work_is_plausible = delegate->DoIdleWork();
     if (run_state->should_quit)
       break;
   }
 }
 
 bool MessagePumpMojo::DoInternalWork(const RunState& run_state, bool block) {
-  const MojoDeadline deadline = block ? GetDeadlineForWait(run_state) : 0;
-  const WaitState wait_state = GetWaitState();
-
-  std::vector<MojoHandleSignalsState> states(wait_state.handles.size());
-  const WaitManyResult wait_many_result =
-      WaitMany(wait_state.handles, wait_state.wait_signals, deadline, &states);
-  const MojoResult result = wait_many_result.result;
-  bool did_work = true;
-  if (result == MOJO_RESULT_OK) {
-    if (wait_many_result.index == 0) {
-      if (states[0].satisfied_signals & MOJO_HANDLE_SIGNAL_PEER_CLOSED) {
-        // The Mojo EDK is shutting down. The ThreadQuitHelper task in
-        // base::Thread won't get run since the control pipe depends on the EDK
-        // staying alive. So quit manually to avoid this thread hanging.
-        Quit();
-      } else {
-        // Control pipe was written to.
-        ReadMessageRaw(read_handle_.get(), NULL, NULL, NULL, NULL,
-                       MOJO_READ_MESSAGE_FLAG_MAY_DISCARD);
-      }
+  bool did_work = block;
+  if (block) {
+    const MojoDeadline deadline = GetDeadlineForWait(run_state);
+    const MojoResult wait_result = Wait(wait_set_handle_.get(),

[sky, 2015/12/11 17:49:01]

How come you don't always call wait?

[Anand Mistry (off Chromium), 2015/12/14 03:16:20]

Added a comment to address.

+                                        MOJO_HANDLE_SIGNAL_READABLE,
+                                        deadline, nullptr);
+    if (wait_result == MOJO_RESULT_OK) {
+      // Handles may be ready. Or not since wake-ups can be spurious in certain
+      // circumstances.
+    } else if (wait_result == MOJO_RESULT_DEADLINE_EXCEEDED) {
+      did_work = false;
     } else {
-      DCHECK(handlers_.find(wait_state.handles[wait_many_result.index]) !=
-             handlers_.end());
-      WillSignalHandler();
-      handlers_[wait_state.handles[wait_many_result.index]]
-          .handler->OnHandleReady(wait_state.handles[wait_many_result.index]);
-      DidSignalHandler();
-    }
-  } else {
-    switch (result) {
-      case MOJO_RESULT_CANCELLED:
-      case MOJO_RESULT_FAILED_PRECONDITION:
-      case MOJO_RESULT_INVALID_ARGUMENT:
-        RemoveInvalidHandle(wait_state, result, wait_many_result.index);
-        break;
-      case MOJO_RESULT_DEADLINE_EXCEEDED:
-        did_work = false;
-        break;
-      default:
-        base::debug::Alias(&result);
-        // Unexpected result is likely fatal, crash so we can determine cause.
-        CHECK(false);
+      base::debug::Alias(&wait_result);
+      // Unexpected result is likely fatal, crash so we can determine cause.
+      CHECK(false);
     }
   }
 
+  const uint32_t kMaxServiced = 8;

[sky, 2015/12/11 17:49:01]

This function is rather lengthy. How about breaking it up? Maybe move this whole
block into its own function?

[Anand Mistry (off Chromium), 2015/12/14 03:16:20]

Done.

+  uint32_t count = kMaxServiced;
+  MojoResult handle_results[kMaxServiced];
+  MojoHandle handles[kMaxServiced] = {MOJO_HANDLE_INVALID};
+
+  const MojoResult get_result = MojoGetReadyHandles(
+      wait_set_handle_.get().value(),
+      &count, handles, handle_results, nullptr);
+  CHECK(get_result == MOJO_RESULT_OK || get_result == MOJO_RESULT_SHOULD_WAIT);
+  if (get_result == MOJO_RESULT_OK) {
+    DCHECK(count);
+    DCHECK_LE(count, kMaxServiced);
+    // Do this in two steps, because notifying a handler may remove/add other
+    // handles that may have also been woken up.
+    // First, enumerate the IDs of the ready handles. Then, iterate over the
+    // handles and only take action if the ID hasn't changed.
+    std::map<Handle, int> ready_handles;
+    for (uint32_t i = 0 ; i < count; i++) {
+      const Handle handle = Handle(handles[i]);
+      // Skip the control handle. It's special.
+      if (handle.value() == read_handle_.get().value())
+        continue;
+      DCHECK(handle.is_valid());
+      const auto it = handlers_.find(handle);
+      // Skip handles that have been removed.

[sky, 2015/12/11 17:49:01]

Under what conditions could this happen? I'm wondering why it isn't a DCHECK?

[Anand Mistry (off Chromium), 2015/12/14 03:16:20]

Added comment to explain why this is possible. It's a consequence of
https://codereview.chromium.org/1475983004, and that handles can be closed
before they are removed from the wait.

+      if (it == handlers_.end())
+        continue;
+      ready_handles[handle] = it->second.id;
+    }
+
+    for (uint32_t i = 0 ; i < count; i++) {
+      const Handle handle = Handle(handles[i]);
+
+      // If the handle has been removed, or it's ID has changed, skip over it.
+      // If the handle's ID has changed, and it still satisfied its signals,
+      // then it'll be caught in the next message pump iteration.
+      const auto it = handlers_.find(handle);
+      if ((handle.value() != read_handle_.get().value()) &&
+          (it == handlers_.end() || it->second.id != ready_handles[handle])) {
+        continue;
+      }
+
+      switch (handle_results[i]) {
+        case MOJO_RESULT_CANCELLED:
+        case MOJO_RESULT_FAILED_PRECONDITION:
+          DVLOG(1) << "Error: " << handle_results[i]
+                   << " handle: " << handle.value();
+          if (handle.value() == read_handle_.get().value()) {
+            // The Mojo EDK is shutting down. The ThreadQuitHelper task in
+            // base::Thread won't get run since the control pipe depends on the
+            // EDK staying alive. So quit manually to avoid this thread hanging.
+            Quit();
+          } else {
+            RemoveInvalidHandle(handle_results[i], handle);
+          }
+          break;
+        case MOJO_RESULT_OK:
+          if (handle.value() == read_handle_.get().value()) {
+            DVLOG(1) << "Signaled control pipe";
+            // Control pipe was written to.
+            ReadMessageRaw(read_handle_.get(), NULL, NULL, NULL, NULL,

[sky, 2015/12/11 17:49:01]

nullptr on these?

[Anand Mistry (off Chromium), 2015/12/14 03:16:20]

Done.

+                           MOJO_READ_MESSAGE_FLAG_MAY_DISCARD);
+          } else {
+            DVLOG(1) << "Handle ready: " << handle.value();
+            SignalHandleReady(handle);
+          }
+          break;
+        default:
+          base::debug::Alias(&i);
+          base::debug::Alias(&handle_results[i]);
+          // Unexpected result is likely fatal, crash so we can determine cause.
+          CHECK(false);
+      }
+    }
+    did_work = true;
+  }
+
   // Notify and remove any handlers whose time has expired. First, iterate over
   // the set of handles that have a deadline, and add the expired handles to a
   // map of <Handle, id>. Then, iterate over those expired handles and remove
   // them. The two-step process is because a handler can add/remove new
   // handlers.
   std::map<Handle, int> expired_handles;
   const base::TimeTicks now(internal::NowTicks());
   for (const Handle handle : deadline_handles_) {
     const auto it = handlers_.find(handle);
     // Expect any handle in |deadline_handles_| to also be in |handlers_| since
@@ skipping 11 matching lines @@
       RemoveHandler(pair.first);
       WillSignalHandler();
       handler->OnHandleError(pair.first, MOJO_RESULT_DEADLINE_EXCEEDED);
       DidSignalHandler();
       did_work = true;
     }
   }
   return did_work;
 }
 
-void MessagePumpMojo::RemoveInvalidHandle(const WaitState& wait_state,
-                                          MojoResult result,
-                                          uint32_t index) {
+void MessagePumpMojo::RemoveInvalidHandle(MojoResult result, Handle handle) {
   // TODO(sky): deal with control pipe going bad.
-  CHECK(result == MOJO_RESULT_INVALID_ARGUMENT ||
-        result == MOJO_RESULT_FAILED_PRECONDITION ||
-        result == MOJO_RESULT_CANCELLED);
-  CHECK_NE(index, 0u);  // Indicates the control pipe went bad.
+  CHECK(result == MOJO_RESULT_FAILED_PRECONDITION ||
+        result == MOJO_RESULT_CANCELLED ||
+        result == MOJO_RESULT_DEADLINE_EXCEEDED);
+  // Indicates the control pipe went bad.
+  CHECK_NE(handle.value(), read_handle_.get().value());
 
   // Remove the handle first, this way if OnHandleError() tries to remove the
   // handle our iterator isn't invalidated.
-  Handle handle = wait_state.handles[index];
-  CHECK(handlers_.find(handle) != handlers_.end());
+  if (handlers_.find(handle) == handlers_.end())

[sky, 2015/12/11 17:49:01]

Under what conditions would this hit?

[Anand Mistry (off Chromium), 2015/12/14 03:16:20]

It did in a previous iteration, but not any more. Put the CHECK back.

+    return;
   MessagePumpMojoHandler* handler = handlers_[handle].handler;
   RemoveHandler(handle);
   WillSignalHandler();
   handler->OnHandleError(handle, result);
   DidSignalHandler();
 }
 
 void MessagePumpMojo::SignalControlPipe() {
   const MojoResult result =
       WriteMessageRaw(write_handle_.get(), NULL, 0, NULL, 0,
                       MOJO_WRITE_MESSAGE_FLAG_NONE);
   if (result == MOJO_RESULT_FAILED_PRECONDITION) {
     // Mojo EDK is shutting down.
     return;
   }
 
   // If we can't write we likely won't wake up the thread and there is a strong
   // chance we'll deadlock.
   CHECK_EQ(MOJO_RESULT_OK, result);
 }
 
-MessagePumpMojo::WaitState MessagePumpMojo::GetWaitState() const {
-  WaitState wait_state;
-  wait_state.handles.push_back(read_handle_.get());
-  wait_state.wait_signals.push_back(
-      MOJO_HANDLE_SIGNAL_READABLE | MOJO_HANDLE_SIGNAL_PEER_CLOSED);
-
-  for (HandleToHandler::const_iterator i = handlers_.begin();
-       i != handlers_.end(); ++i) {
-    wait_state.handles.push_back(i->first);
-    wait_state.wait_signals.push_back(i->second.wait_signals);
-  }
-  return wait_state;
-}
-
 MojoDeadline MessagePumpMojo::GetDeadlineForWait(
     const RunState& run_state) const {
   const base::TimeTicks now(internal::NowTicks());
   MojoDeadline deadline = TimeTicksToMojoDeadline(run_state.delayed_work_time,
                                                   now);
   for (const Handle handle : deadline_handles_) {
     auto it = handlers_.find(handle);
     DCHECK(it != handlers_.end());
     deadline = std::min(
         TimeTicksToMojoDeadline(it->second.deadline, now), deadline);
   }
   return deadline;
 }
 
+void MessagePumpMojo::SignalHandleReady(Handle handle) {
+  DCHECK(handlers_.find(handle) != handlers_.end());
+  WillSignalHandler();
+  handlers_[handle].handler->OnHandleReady(handle);
+  DidSignalHandler();
+}
+
 void MessagePumpMojo::WillSignalHandler() {
   FOR_EACH_OBSERVER(Observer, observers_, WillSignalHandler());
 }
 
 void MessagePumpMojo::DidSignalHandler() {
   FOR_EACH_OBSERVER(Observer, observers_, DidSignalHandler());
 }
 
 }  // namespace common
 }  // namespace mojo