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"

[sky, 2015/12/08 18:26:14]

It's a bit hard to review this patch without knowing waitset. Can you ping it
once waitset lands?

[Anand Mistry (off Chromium), 2015/12/09 00:01:21]

Will do.

The API is defined in https://codereview.chromium.org/1429553002/, for anyone
that casually comes across this CL.

 
 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 asynchrnous, it's possible for the
+  // handle to no longer be open at this point.
+  CHECK(result == MOJO_RESULT_OK || result == MOJO_RESULT_INVALID_ARGUMENT);
 }
 
 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();
+  bool did_work = block;
+  if (block) {
+    const MojoDeadline deadline = block ? GetDeadlineForWait(run_state) : 0;
 
-  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);
-      }
+
+    const MojoResult wait_result = Wait(wait_set_handle_.get(),
+                                        MOJO_HANDLE_SIGNAL_READABLE,
+                                        deadline, nullptr);
+    DVLOG(1) << "Wait result: " << wait_result << ", block: " << block;
+    if (wait_result == MOJO_RESULT_OK) {
+    } else if (wait_result == MOJO_RESULT_DEADLINE_EXCEEDED) {
+      DVLOG(1) << "Deadline expired, block: " << block;
+      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;
+  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 taking 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;
+      CHECK(handle.is_valid());
+      const auto it = handlers_.find(handle);
+      // Skip handles that have been removed.
+      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;
+      }
+
+      CHECK(handle.is_valid());
+      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,
+                           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;
+  } else {
+    DVLOG(1) << "No handles ready, block: " << block;
+  }
+
   // 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())
+    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) {
+  CHECK(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