Mojo EDK: Introduce MojoQueryHandleSignalsState API

mojo/edk/test/mojo_test_base.cc

 // Copyright 2016 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/edk/test/mojo_test_base.h"
 
 #include "base/memory/ptr_util.h"
+#include "base/memory/ref_counted.h"
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.h"
+#include "base/synchronization/waitable_event.h"
 #include "mojo/edk/embedder/embedder.h"
 #include "mojo/edk/system/handle_signals_state.h"
 #include "mojo/public/c/system/buffer.h"
 #include "mojo/public/c/system/data_pipe.h"
 #include "mojo/public/c/system/functions.h"
+#include "mojo/public/c/system/watcher.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 #if defined(OS_MACOSX) && !defined(OS_IOS)
 #include "base/mac/mach_port_broker.h"
 #endif
 
 namespace mojo {
 namespace edk {
 namespace test {
 
+namespace {
+
+class Waiter {
+ public:
+  Waiter() {}
+  ~Waiter() {}
+
+  MojoResult Wait(MojoHandle handle,
+                  MojoHandleSignals signals,
+                  MojoHandleSignalsState* state) {
+    MojoHandle watcher;
+    MojoCreateWatcher(&Context::OnNotification, &watcher);
+
+    context_ = new Context();
+
+    // Balanced by OnNotification in the |MOJO_RESULT_CANCELLED| case.
+    context_->AddRef();
+
+    MojoResult rv = MojoWatch(watcher, handle, signals,
+                              reinterpret_cast<uintptr_t>(context_.get()));
+    DCHECK_EQ(MOJO_RESULT_OK, rv);
+
+    uint32_t num_ready_contexts = 1;
+    uintptr_t ready_context;
+    MojoResult ready_result;
+    MojoHandleSignalsState ready_state;
+    rv = MojoArmWatcher(watcher, &num_ready_contexts, &ready_context,
+                        &ready_result, &ready_state);
+    if (rv == MOJO_RESULT_FAILED_PRECONDITION) {
+      MojoClose(watcher);
+      DCHECK_EQ(1u, num_ready_contexts);
+      if (state)
+        *state = ready_state;
+      return ready_result;
+    }
+
+    // Wait for the first notification.
+    context_->event().Wait();
+
+    ready_result = context_->wait_result();
+    DCHECK_NE(MOJO_RESULT_UNKNOWN, ready_result);
+
+    if (state)
+      *state = context_->wait_state();
+
+    MojoClose(watcher);
+
+    return ready_result;
+  }
+
+ private:
+  class Context : public base::RefCountedThreadSafe<Context> {
+   public:
+    Context()
+        : event_(base::WaitableEvent::ResetPolicy::AUTOMATIC,
+                 base::WaitableEvent::InitialState::NOT_SIGNALED) {}
+
+    base::WaitableEvent& event() { return event_; }
+    MojoResult wait_result() const { return wait_result_; }
+    MojoHandleSignalsState wait_state() const { return wait_state_; }
+
+    static void OnNotification(uintptr_t context_value,
+                               MojoResult result,
+                               MojoHandleSignalsState state,
+                               MojoWatcherNotificationFlags flags) {
+      auto* context = reinterpret_cast<Context*>(context_value);
+      context->Notify(result, state);
+      if (result == MOJO_RESULT_CANCELLED)
+        context->Release();
+    }
+
+   private:
+    friend class base::RefCountedThreadSafe<Context>;
+
+    ~Context() {}
+
+    void Notify(MojoResult result, MojoHandleSignalsState state) {
+      if (wait_result_ == MOJO_RESULT_UNKNOWN) {
+        wait_result_ = result;
+        wait_state_ = state;
+      }
+      event_.Signal();
+    }
+
+    base::WaitableEvent event_;
+    MojoResult wait_result_ = MOJO_RESULT_UNKNOWN;
+    MojoHandleSignalsState wait_state_ = {0, 0};
+
+    DISALLOW_COPY_AND_ASSIGN(Context);
+  };
+
+  scoped_refptr<Context> context_;
+
+  DISALLOW_COPY_AND_ASSIGN(Waiter);
+};
+
+}  // namespace
 
 #if defined(OS_MACOSX) && !defined(OS_IOS)
 namespace {
 base::MachPortBroker* g_mach_broker = nullptr;
 }
 #endif
 
 MojoTestBase::MojoTestBase() {
 #if defined(OS_MACOSX) && !defined(OS_IOS)
   if (!g_mach_broker) {
@@ skipping 88 matching lines @@
 // static
 void MojoTestBase::WriteMessage(MojoHandle mp, const std::string& message) {
   WriteMessageWithHandles(mp, message, nullptr, 0);
 }
 
 // static
 std::string MojoTestBase::ReadMessageWithHandles(
     MojoHandle mp,
     MojoHandle* handles,
     uint32_t expected_num_handles) {
-  CHECK_EQ(MojoWait(mp, MOJO_HANDLE_SIGNAL_READABLE, MOJO_DEADLINE_INDEFINITE,
-                    nullptr),
-           MOJO_RESULT_OK);
+  CHECK_EQ(WaitForSignals(mp, MOJO_HANDLE_SIGNAL_READABLE), MOJO_RESULT_OK);
 
   uint32_t message_size = 0;
   uint32_t num_handles = 0;
   CHECK_EQ(MojoReadMessage(mp, nullptr, &message_size, nullptr, &num_handles,
                            MOJO_READ_MESSAGE_FLAG_NONE),
            MOJO_RESULT_RESOURCE_EXHAUSTED);
   CHECK_EQ(expected_num_handles, num_handles);
 
   std::string message(message_size, 'x');
   CHECK_EQ(MojoReadMessage(mp, &message[0], &message_size, handles,
                            &num_handles, MOJO_READ_MESSAGE_FLAG_NONE),
            MOJO_RESULT_OK);
   CHECK_EQ(message_size, message.size());
   CHECK_EQ(num_handles, expected_num_handles);
 
   return message;
 }
 
 // static
 std::string MojoTestBase::ReadMessageWithOptionalHandle(MojoHandle mp,
                                                         MojoHandle* handle) {
-  CHECK_EQ(MojoWait(mp, MOJO_HANDLE_SIGNAL_READABLE, MOJO_DEADLINE_INDEFINITE,
-                    nullptr),
-           MOJO_RESULT_OK);
+  CHECK_EQ(WaitForSignals(mp, MOJO_HANDLE_SIGNAL_READABLE), MOJO_RESULT_OK);
 
   uint32_t message_size = 0;
   uint32_t num_handles = 0;
   CHECK_EQ(MojoReadMessage(mp, nullptr, &message_size, nullptr, &num_handles,
                            MOJO_READ_MESSAGE_FLAG_NONE),
            MOJO_RESULT_RESOURCE_EXHAUSTED);
   CHECK(num_handles == 0 || num_handles == 1);
 
   CHECK(handle);
 
@@ skipping 14 matching lines @@
 
 // static
 std::string MojoTestBase::ReadMessage(MojoHandle mp) {
   return ReadMessageWithHandles(mp, nullptr, 0);
 }
 
 // static
 void MojoTestBase::ReadMessage(MojoHandle mp,
                                char* data,
                                size_t num_bytes) {
-  CHECK_EQ(MojoWait(mp, MOJO_HANDLE_SIGNAL_READABLE, MOJO_DEADLINE_INDEFINITE,
-                    nullptr),
-           MOJO_RESULT_OK);
+  CHECK_EQ(WaitForSignals(mp, MOJO_HANDLE_SIGNAL_READABLE), MOJO_RESULT_OK);
 
   uint32_t message_size = 0;
   uint32_t num_handles = 0;
   CHECK_EQ(MojoReadMessage(mp, nullptr, &message_size, nullptr, &num_handles,
                            MOJO_READ_MESSAGE_FLAG_NONE),
            MOJO_RESULT_RESOURCE_EXHAUSTED);
   CHECK_EQ(num_handles, 0u);
   CHECK_EQ(message_size, num_bytes);
 
   CHECK_EQ(MojoReadMessage(mp, data, &message_size, nullptr, &num_handles,
@@ skipping 74 matching lines @@
   options.element_num_bytes = 1;
   options.capacity_num_bytes = static_cast<uint32_t>(capacity);
 
   MojoCreateDataPipe(&options, p0, p1);
   CHECK_NE(*p0, MOJO_HANDLE_INVALID);
   CHECK_NE(*p1, MOJO_HANDLE_INVALID);
 }
 
 // static
 void MojoTestBase::WriteData(MojoHandle producer, const std::string& data) {
-  CHECK_EQ(MojoWait(producer, MOJO_HANDLE_SIGNAL_WRITABLE,
-                    MOJO_DEADLINE_INDEFINITE, nullptr),
+  CHECK_EQ(WaitForSignals(producer, MOJO_HANDLE_SIGNAL_WRITABLE),
            MOJO_RESULT_OK);
   uint32_t num_bytes = static_cast<uint32_t>(data.size());
   CHECK_EQ(MojoWriteData(producer, data.data(), &num_bytes,
                          MOJO_WRITE_DATA_FLAG_ALL_OR_NONE),
            MOJO_RESULT_OK);
   CHECK_EQ(num_bytes, static_cast<uint32_t>(data.size()));
 }
 
 // static
 std::string MojoTestBase::ReadData(MojoHandle consumer, size_t size) {
-  CHECK_EQ(MojoWait(consumer, MOJO_HANDLE_SIGNAL_READABLE,
-                    MOJO_DEADLINE_INDEFINITE, nullptr),
+  CHECK_EQ(WaitForSignals(consumer, MOJO_HANDLE_SIGNAL_READABLE),
            MOJO_RESULT_OK);
   std::vector<char> buffer(size);
   uint32_t num_bytes = static_cast<uint32_t>(size);
   CHECK_EQ(MojoReadData(consumer, buffer.data(), &num_bytes,
                          MOJO_WRITE_DATA_FLAG_ALL_OR_NONE),
            MOJO_RESULT_OK);
   CHECK_EQ(num_bytes, static_cast<uint32_t>(size));
 
   return std::string(buffer.data(), buffer.size());
 }
 
+// static
+MojoHandleSignalsState MojoTestBase::GetSignalsState(MojoHandle handle) {
+  MojoHandleSignalsState signals_state;
+  CHECK_EQ(MOJO_RESULT_OK, MojoQueryHandleSignalsState(handle, &signals_state));
+  return signals_state;
+}
+
+// static
+MojoResult MojoTestBase::WaitForSignals(MojoHandle handle,
+                                        MojoHandleSignals signals,
+                                        MojoHandleSignalsState* state) {
+  Waiter waiter;
+  return waiter.Wait(handle, signals, state);
+}
+
 }  // namespace test
 }  // namespace edk
 }  // namespace mojo