Mojo: Armed Watchers

content/common/message_port.cc

Index: content/common/message_port.cc
diff --git a/content/common/message_port.cc b/content/common/message_port.cc
index c89d7b5a0d0cc313290329602fedc13e6dff9150..b7a77645d1f2faf7f32d9de2c39a74a0880b18a1 100644
--- a/content/common/message_port.cc
+++ b/content/common/message_port.cc
@@ -4,7 +4,9 @@
 
 #include "content/common/message_port.h"
 
+#include "base/bind.h"
 #include "base/logging.h"
+#include "base/threading/thread_task_runner_handle.h"
 
 namespace content {
 
@@ -31,7 +33,7 @@ const mojo::ScopedMessagePipeHandle& MessagePort::GetHandle() const {
 }
 
 mojo::ScopedMessagePipeHandle MessagePort::ReleaseHandle() const {
-  state_->CancelWatch();
+  state_->UnregisterWatcher();
   return std::move(state_->handle_);
 }
 
@@ -126,13 +128,13 @@ bool MessagePort::GetMessage(base::string16* encoded_message,
 }
 
 void MessagePort::SetCallback(const base::Closure& callback) {
-  state_->CancelWatch();
+  state_->UnregisterWatcher();
   state_->callback_ = callback;
-  state_->AddWatch();
+  state_->RegisterWatcher();
 }
 
 void MessagePort::ClearCallback() {
-  state_->CancelWatch();
+  state_->UnregisterWatcher();
   state_->callback_.Reset();
 }
 
@@ -143,44 +145,77 @@ MessagePort::State::State(mojo::ScopedMessagePipeHandle handle)
     : handle_(std::move(handle)) {
 }
 
-void MessagePort::State::AddWatch() {
+void MessagePort::State::RegisterWatcher() {
   if (!callback_)
     return;
 
   // NOTE: An HTML MessagePort does not receive an event to tell it when the
   // peer has gone away, so we only watch for readability here.
-  MojoResult rv = MojoWatch(handle_.get().value(),
-                            MOJO_HANDLE_SIGNAL_READABLE,
-                            &MessagePort::State::OnHandleReady,
-                            reinterpret_cast<uintptr_t>(this));
-  if (rv != MOJO_RESULT_OK)
-    DVLOG(1) << this << " MojoWatch failed: " << rv;
+  MojoResult rv =
+      MojoRegisterWatcher(handle_.get().value(), MOJO_HANDLE_SIGNAL_READABLE,
+                          &MessagePort::State::CallOnHandleReady,
+                          reinterpret_cast<uintptr_t>(this));
+  if (rv != MOJO_RESULT_OK) {
+    DVLOG(1) << this << " MojoRegisterWatcher failed: " << rv;
+    return;
+  }
+
+  ArmWatcher();
 }
 
-void MessagePort::State::CancelWatch() {
+void MessagePort::State::UnregisterWatcher() {
   if (!callback_)
     return;
 
   // NOTE: This synchronizes with the thread where OnHandleReady runs so we are
   // sure to not be racing with it.
-  MojoCancelWatch(handle_.get().value(), reinterpret_cast<uintptr_t>(this));
+  MojoUnregisterWatcher(handle_.get().value(),
+                        reinterpret_cast<uintptr_t>(this));
 }
 
-// static
-void MessagePort::State::OnHandleReady(
-    uintptr_t context,
-    MojoResult result,
-    MojoHandleSignalsState signals_state,
-    MojoWatchNotificationFlags flags) {
+MessagePort::State::~State() {
+  UnregisterWatcher();
+}
+
+void MessagePort::State::ArmWatcher() {
+  if (!callback_)
+    return;
+
+  MojoResult rv =
+      MojoArmWatcher(handle_.get().value(), reinterpret_cast<uintptr_t>(this));
+  if (rv == MOJO_RESULT_OK)
+    return;
+
+  if (rv == MOJO_RESULT_ALREADY_EXISTS) {
+    // The handle is already signalled, so we trigger a callback immediately.
+    base::ThreadTaskRunnerHandle::Get()->PostTask(
+        FROM_HERE, base::Bind(&State::OnHandleReady, this, MOJO_RESULT_OK));
+    return;
+  }
+
+  if (rv == MOJO_RESULT_FAILED_PRECONDITION) {
+    DVLOG(1) << this << " MojoArmWatcher failed because of a broken pipe.";
+    return;
+  }
+
+  NOTREACHED();
+}
+
+void MessagePort::State::OnHandleReady(MojoResult result) {
   if (result == MOJO_RESULT_OK) {
-    reinterpret_cast<MessagePort::State*>(context)->callback_.Run();
+    callback_.Run();
+    ArmWatcher();
   } else {
     // And now his watch is ended.
   }
 }
 
-MessagePort::State::~State() {
-  CancelWatch();
+// static
+void MessagePort::State::CallOnHandleReady(uintptr_t context,
+                                           MojoResult result,
+                                           MojoHandleSignalsState signals_state,
+                                           MojoWatchNotificationFlags flags) {
+  reinterpret_cast<State*>(context)->OnHandleReady(result);
 }
 
 }  // namespace content