[mojo-edk] Prevent AcceptIncomingMessages flushing aggressively

mojo/edk/system/watch_unittest.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 <functional>
 
 #include "base/macros.h"
+#include "base/memory/ref_counted.h"
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.h"
+#include "base/single_thread_task_runner.h"
+#include "base/threading/thread_task_runner_handle.h"
 #include "mojo/edk/system/request_context.h"
 #include "mojo/edk/test/mojo_test_base.h"
 #include "mojo/public/c/system/functions.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace mojo {
 namespace edk {
 namespace {
 
 void IgnoreResult(uintptr_t context,
                   MojoResult result,
                   MojoHandleSignalsState signals,
                   MojoWatchNotificationFlags flags) {
 }
 
 // A test helper class for watching a handle. The WatchHelper instance is used
 // as a watch context for a single watch callback.
 class WatchHelper {
  public:
   using Callback =
       std::function<void(MojoResult result, MojoHandleSignalsState state)>;
 
-  WatchHelper() {}
+  WatchHelper() : task_runner_(base::ThreadTaskRunnerHandle::Get()) {}
   ~WatchHelper() {
     CHECK(!watching_);
   }
 
   void Watch(MojoHandle handle,
              MojoHandleSignals signals,
              const Callback& callback) {
     CHECK(!watching_);
 
     handle_ = handle;
     callback_ = callback;
     watching_ = true;
     CHECK_EQ(MOJO_RESULT_OK, MojoWatch(handle_, signals, &WatchHelper::OnNotify,
                                        reinterpret_cast<uintptr_t>(this)));
   }
 
   bool is_watching() const { return watching_; }
 
   void Cancel() {
     CHECK_EQ(MOJO_RESULT_OK,
              MojoCancelWatch(handle_, reinterpret_cast<uintptr_t>(this)));
     CHECK(watching_);
     watching_ = false;
   }
 
-  void ExpectSystemNotifications() { expect_system_notifications_ = true; }
-
  private:
   static void OnNotify(uintptr_t context,
                        MojoResult result,
                        MojoHandleSignalsState state,
                        MojoWatchNotificationFlags flags) {
     WatchHelper* watcher = reinterpret_cast<WatchHelper*>(context);
+    watcher->task_runner_->PostTask(
+        FROM_HERE,
+        base::Bind(&NotifyOnMainThread, context, result, state, flags));
+  }
+
+  static void NotifyOnMainThread(uintptr_t context,
+                                 MojoResult result,
+                                 MojoHandleSignalsState state,
+                                 MojoWatchNotificationFlags flags) {
+    WatchHelper* watcher = reinterpret_cast<WatchHelper*>(context);
     CHECK(watcher->watching_);
     if (result == MOJO_RESULT_CANCELLED)
       watcher->watching_ = false;
-    CHECK_EQ(flags, watcher->expect_system_notifications_?
-        MOJO_WATCH_NOTIFICATION_FLAG_FROM_SYSTEM :
-        MOJO_WATCH_NOTIFICATION_FLAG_NONE);
     watcher->callback_(result, state);
   }
 
+  scoped_refptr<base::SingleThreadTaskRunner> task_runner_;
   bool watching_ = false;
   MojoHandle handle_;
   Callback callback_;
-  bool expect_system_notifications_ = false;
 
   DISALLOW_COPY_AND_ASSIGN(WatchHelper);
 };
 
 class WatchTest : public test::MojoTestBase {
  public:
   WatchTest() {}
   ~WatchTest() override {}
 
  protected:
@@ skipping 153 matching lines @@
 
 TEST_F(WatchTest, WatchWhileSatisfied) {
   MojoHandle a, b;
   CreateMessagePipe(&a, &b);
 
   // Write to |a| and then start watching |b|. The callback should be invoked
   // synchronously.
   WriteMessage(a, "hey");
   bool signaled = false;
   WatchHelper b_watcher;
+  base::RunLoop loop;
   b_watcher.Watch(
       b, MOJO_HANDLE_SIGNAL_READABLE,
       [&] (MojoResult result, MojoHandleSignalsState state) {
         EXPECT_EQ(MOJO_RESULT_OK, result);
         EXPECT_EQ(MOJO_HANDLE_SIGNAL_READABLE,
                   state.satisfied_signals & MOJO_HANDLE_SIGNAL_READABLE);
         signaled = true;
+        loop.Quit();
       });
+  loop.Run();
   EXPECT_TRUE(signaled);
   b_watcher.Cancel();
 
   CloseHandle(a);
   CloseHandle(b);
 }
 
 TEST_F(WatchTest, WatchWhileUnsatisfiable) {
   MojoHandle a, b;
   CreateMessagePipe(&a, &b);
@@ skipping 142 matching lines @@
 
 TEST_F(WatchTest, NestedCancellation) {
   // Verifies that cancellations in nested system request contexts preempt
   // other notifications for the same watcher. This tests against the condition
   // hit by http://crbug.com/622298.
 
   MojoHandle a, b, c, d;
   CreateMessagePipe(&a, &b);
   CreateMessagePipe(&c, &d);
 
+  base::RunLoop loop;
   bool a_watcher_run = false;
   WatchHelper a_watcher;
   a_watcher.Watch(
       a, MOJO_HANDLE_SIGNAL_READABLE,
       [&](MojoResult result, MojoHandleSignalsState state) {
         a_watcher_run = true;
       });
 
   WatchHelper c_watcher;
   c_watcher.Watch(
       c, MOJO_HANDLE_SIGNAL_READABLE,
       [&](MojoResult result, MojoHandleSignalsState state) {
         // This will trigger a notification on |a_watcher| above to be executed
         // once this handler finishes running...
         CloseHandle(b);
 
         // ...but this should prevent that notification from dispatching because
         // |a_watcher| is now cancelled.
         a_watcher.Cancel();
+
+        loop.Quit();
       });
-  c_watcher.ExpectSystemNotifications();
 
   {
     // Force "system" notifications for the synchronous behavior required to
     // test this case.
     mojo::edk::RequestContext request_context(
         mojo::edk::RequestContext::Source::SYSTEM);
 
     // Trigger the |c_watcher| callback above.
     CloseHandle(d);
   }
 
+  loop.Run();
+
   EXPECT_FALSE(a_watcher.is_watching());
   EXPECT_FALSE(a_watcher_run);
 
   c_watcher.Cancel();
 
   CloseHandle(a);
   CloseHandle(c);
 }
 
 }  // namespace
 }  // namespace edk
 }  // namespace mojo