Convert SyncHandleRegistry to use a CallbackList.

This changes the API to look like CallbackList's: instead of calling
UnregisterEvent(), callers destroy a subscription object they get from
registering.  This is a bit easier for callers to deal with.  However,
the necessity of also synchronously clearing the event from the WaitSet
makes the implementation a bit irritating, since while CallbackList does
provide a hook to be called back on callback removal, that happens after
notification is finished, rather than right at subscription destruction.
Instead of using the hook, wrap the CallbackList::Subscriptions in an
object that will also remove the WaitSet item if need be.

Bug: none
Change-Id: Ideb74b1435bfaad1647324e5ba5f90976c268e2e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2343956
Commit-Queue: Peter Kasting <pkasting@chromium.org>
Reviewed-by: Ken Rockot <rockot@google.com>
Cr-Commit-Position: refs/heads/master@{#796728}

ipc/ipc_sync_channel.cc

@@ -667,28 +667,27 @@ void SyncChannel::WaitForReply(mojo::SyncHandleRegistry* registry,
 
   while (true) {
     bool dispatch = false;
-    bool send_done = false;
     bool should_pump_messages = false;
-    base::RepeatingClosure on_send_done_callback =
-        base::BindRepeating(&OnEventReady, &send_done);
-    registry->RegisterEvent(context->GetSendDoneEvent(), on_send_done_callback);
+    {
+      bool send_done = false;
+      mojo::SyncHandleRegistry::EventCallbackSubscription
+          send_done_subscription = registry->RegisterEvent(
+              context->GetSendDoneEvent(),
+              base::BindRepeating(&OnEventReady, &send_done));
 
-    base::RepeatingClosure on_pump_messages_callback;
+      mojo::SyncHandleRegistry::EventCallbackSubscription
+          pump_messages_subsciption;
       if (pump_messages_event) {
-      on_pump_messages_callback =
-          base::BindRepeating(&OnEventReady, &should_pump_messages);
-      registry->RegisterEvent(pump_messages_event, on_pump_messages_callback);
+        pump_messages_subsciption = registry->RegisterEvent(
+            pump_messages_event,
+            base::BindRepeating(&OnEventReady, &should_pump_messages));
       }
 
-    const bool* stop_flags[] = { &dispatch, &send_done, &should_pump_messages };
+      const bool* stop_flags[] = {&dispatch, &send_done, &should_pump_messages};
       context->received_sync_msgs()->BlockDispatch(&dispatch);
       registry->Wait(stop_flags, 3);
       context->received_sync_msgs()->UnblockDispatch();
-
-    registry->UnregisterEvent(context->GetSendDoneEvent(),
-                              on_send_done_callback);
-    if (pump_messages_event)
-      registry->UnregisterEvent(pump_messages_event, on_pump_messages_callback);
+    }
 
     if (dispatch) {
       // We're waiting for a reply, but we received a blocking synchronous call.

ipc/ipc_sync_message_filter.cc

@@ -69,24 +69,25 @@ bool SyncMessageFilter::Send(Message* message) {
     }
   }
 
+  {
     bool done = false;
     bool shutdown = false;
     scoped_refptr<mojo::SyncHandleRegistry> registry =
         mojo::SyncHandleRegistry::current();
-  auto on_shutdown_callback = base::BindRepeating(&OnEventReady, &shutdown);
-  auto on_done_callback = base::BindRepeating(&OnEventReady, &done);
-  registry->RegisterEvent(shutdown_event_, on_shutdown_callback);
-  registry->RegisterEvent(&done_event, on_done_callback);
-
-  const bool* stop_flags[] = { &done, &shutdown };
+    mojo::SyncHandleRegistry::EventCallbackSubscription shutdown_subscription =
+        registry->RegisterEvent(shutdown_event_,
+                                base::BindRepeating(&OnEventReady, &shutdown));
+    mojo::SyncHandleRegistry::EventCallbackSubscription done_subscription =
+        registry->RegisterEvent(&done_event,
+                                base::BindRepeating(&OnEventReady, &done));
+
+    const bool* stop_flags[] = {&done, &shutdown};
     registry->Wait(stop_flags, 2);
     if (done) {
       TRACE_EVENT_FLOW_END0("toplevel.flow", "SyncMessageFilter::Send",
                             &done_event);
     }
-
-  registry->UnregisterEvent(shutdown_event_, on_shutdown_callback);
-  registry->UnregisterEvent(&done_event, on_done_callback);
+  }
 
   {
     base::AutoLock auto_lock(lock_);

mojo/public/cpp/bindings/lib/sync_event_watcher.cc

@@ -20,8 +20,6 @@ SyncEventWatcher::SyncEventWatcher(base::WaitableEvent* event,
 
 SyncEventWatcher::~SyncEventWatcher() {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
-  if (registered_)
-    registry_->UnregisterEvent(event_, callback_);
   destroyed_->data = true;
 }
 
@@ -34,10 +32,6 @@ bool SyncEventWatcher::SyncWatch(const bool** stop_flags,
                                  size_t num_stop_flags) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   IncrementRegisterCount();
-  if (!registered_) {
-    DecrementRegisterCount();
-    return false;
-  }
 
   // This object may be destroyed during the Wait() call. So we have to preserve
   // the boolean that Wait uses.
@@ -60,20 +54,14 @@ bool SyncEventWatcher::SyncWatch(const bool** stop_flags,
 }
 
 void SyncEventWatcher::IncrementRegisterCount() {
-  register_request_count_++;
-  if (!registered_) {
-    registry_->RegisterEvent(event_, callback_);
-    registered_ = true;
-  }
+  if (register_request_count_++ == 0)
+    subscription_ = registry_->RegisterEvent(event_, callback_);
 }
 
 void SyncEventWatcher::DecrementRegisterCount() {
   DCHECK_GT(register_request_count_, 0u);
-  register_request_count_--;
-  if (register_request_count_ == 0 && registered_) {
-    registry_->UnregisterEvent(event_, callback_);
-    registered_ = false;
-  }
+  if (--register_request_count_ == 0)
+    subscription_.reset();
 }
 
 }  // namespace mojo

mojo/public/cpp/bindings/lib/sync_handle_registry.cc

@@ -5,7 +5,9 @@
 #include "mojo/public/cpp/bindings/sync_handle_registry.h"
 
 #include <algorithm>
+#include <utility>
 
+#include "base/auto_reset.h"
 #include "base/check_op.h"
 #include "base/no_destructor.h"
 #include "base/stl_util.h"
@@ -15,6 +17,19 @@
 
 namespace mojo {
 
+SyncHandleRegistry::Subscription::Subscription(base::OnceClosure remove_closure,
+                                               EventCallbackList* callbacks,
+                                               EventCallback event_callback)
+    : remove_runner_(std::move(remove_closure)),
+      subscription_(callbacks->Add(std::move(event_callback))) {}
+
+SyncHandleRegistry::Subscription::Subscription(Subscription&&) = default;
+
+SyncHandleRegistry::Subscription& SyncHandleRegistry::Subscription::operator=(
+    Subscription&&) = default;
+
+SyncHandleRegistry::Subscription::~Subscription() = default;
+
 // static
 scoped_refptr<SyncHandleRegistry> SyncHandleRegistry::current() {
   static base::NoDestructor<
@@ -57,11 +72,12 @@ void SyncHandleRegistry::UnregisterHandle(const Handle& handle) {
   handles_.erase(handle);
 }
 
-void SyncHandleRegistry::RegisterEvent(base::WaitableEvent* event,
-                                       base::RepeatingClosure callback) {
+SyncHandleRegistry::EventCallbackSubscription SyncHandleRegistry::RegisterEvent(
+    base::WaitableEvent* event,
+    EventCallback callback) {
   auto it = events_.find(event);
   if (it == events_.end()) {
-    auto result = events_.emplace(event, EventCallbackList{});
+    auto result = events_.emplace(event, std::make_unique<EventCallbackList>());
     it = result.first;
   }
 
@@ -70,43 +86,29 @@ void SyncHandleRegistry::RegisterEvent(base::WaitableEvent* event,
   // callbacks to see if any are valid.
   wait_set_.AddEvent(event);
 
-  it->second.container().push_back(std::move(callback));
-}
-
-void SyncHandleRegistry::UnregisterEvent(base::WaitableEvent* event,
-                                         base::RepeatingClosure callback) {
-  auto it = events_.find(event);
-  if (it == events_.end())
-    return;
-
-  bool has_valid_callbacks = false;
-  auto& callbacks = it->second.container();
-  if (is_dispatching_event_callbacks_) {
-    // Not safe to remove any elements from |callbacks| here since an outer
-    // stack frame is currently iterating over it in Wait().
-    for (auto& cb : callbacks) {
-      if (cb == callback)
-        cb.Reset();
-      else if (cb)
-        has_valid_callbacks = true;
-    }
-    remove_invalid_event_callbacks_after_dispatch_ = true;
-  } else {
-    callbacks.erase(std::remove(callbacks.begin(), callbacks.end(), callback),
-                    callbacks.end());
-    if (callbacks.empty())
-      events_.erase(it);
-    else
-      has_valid_callbacks = true;
-  }
-
-  if (!has_valid_callbacks) {
-    // Regardless of whether or not we're nested within a Wait(), we need to
-    // ensure that |event| is removed from the WaitSet before returning if this
-    // was the last callback registered for it.
-    MojoResult rv = wait_set_.RemoveEvent(event);
+  // Return an object that will synchronously clear the entry for |event| when
+  // its last callback is destroyed.
+  const auto remove_closure = [](EventCallbackList* callbacks,
+                                 WaitSet* wait_set,
+                                 base::WaitableEvent* event) {
+    // |callbacks| is guaranteed to be valid here. The callbacks are repeating
+    // and are thus only removed by their subscriptions being destroyed, so it's
+    // impossible for empty() to be true until the last subscription has been
+    // destroyed.  Since Wait() only deletes a callback list once it's empty,
+    // and this callback runs synchronously with subscription destruction, it's
+    // impossible for |callbacks| to be deleted before this gets to run at the
+    // destruction of the last remaining subscription.
+    if (callbacks->empty()) {
+      // If this was the last callback registered for |event|, ensure that it's
+      // removed from the WaitSet before returning. Otherwise a nested Wait()
+      // call inside the scope that destroys the subscription will fail.
+      const MojoResult rv = wait_set->RemoveEvent(event);
       DCHECK_EQ(MOJO_RESULT_OK, rv);
     }
+  };
+  return std::make_unique<Subscription>(
+      base::BindOnce(remove_closure, it->second.get(), &wait_set_, event),
+      it->second.get(), std::move(callback));
 }
 
 bool SyncHandleRegistry::Wait(const bool* should_stop[], size_t count) {
@@ -138,28 +140,18 @@ bool SyncHandleRegistry::Wait(const bool* should_stop[], size_t count) {
     if (ready_event) {
       const auto iter = events_.find(ready_event);
       DCHECK(iter != events_.end());
-      bool was_dispatching_event_callbacks = is_dispatching_event_callbacks_;
-      is_dispatching_event_callbacks_ = true;
-
-      // NOTE: It's possible for the container to be extended by any of these
-      // callbacks if they call RegisterEvent, so we are careful to iterate by
-      // index. Also note that conversely, elements cannot be *removed* from the
-      // container, by any of these callbacks, so it is safe to assume the size
-      // only stays the same or increases, with no elements changing position.
-      auto& callbacks = iter->second.container();
-      for (size_t i = 0; i < callbacks.size(); ++i) {
-        auto& callback = callbacks[i];
-        if (callback)
-          callback.Run();
+
+      {
+        base::AutoReset<bool> in_nested_wait(&in_nested_wait_, true);
+        iter->second->Notify();
       }
 
-      is_dispatching_event_callbacks_ = was_dispatching_event_callbacks;
-      if (!was_dispatching_event_callbacks &&
-          remove_invalid_event_callbacks_after_dispatch_) {
-        // If we've had events unregistered within any callback dispatch, now is
-        // a good time to prune them from the map.
-        RemoveInvalidEventCallbacks();
-        remove_invalid_event_callbacks_after_dispatch_ = false;
+      // Notify() above may have both added and removed event registrations, for
+      // any event.  If we're in the outermost frame, prune any empty map
+      // entries to avoid unbounded growth.
+      if (!in_nested_wait_) {
+        base::EraseIf(events_,
+                      [](const auto& entry) { return entry.second->empty(); });
       }
     }
   };
@@ -171,19 +163,4 @@ SyncHandleRegistry::SyncHandleRegistry() = default;
 
 SyncHandleRegistry::~SyncHandleRegistry() = default;
 
-void SyncHandleRegistry::RemoveInvalidEventCallbacks() {
-  for (auto it = events_.begin(); it != events_.end();) {
-    auto& callbacks = it->second.container();
-    callbacks.erase(std::remove_if(callbacks.begin(), callbacks.end(),
-                                   [](const base::RepeatingClosure& callback) {
-                                     return !callback;
-                                   }),
-                    callbacks.end());
-    if (callbacks.empty())
-      events_.erase(it++);
-    else
-      ++it;
-  }
-}
-
 }  // namespace mojo

mojo/public/cpp/bindings/sync_event_watcher.h

@@ -53,8 +53,7 @@ class COMPONENT_EXPORT(MOJO_CPP_BINDINGS) SyncEventWatcher {
   base::WaitableEvent* const event_;
   const base::RepeatingClosure callback_;
 
-  // Whether |event_| has been registered with SyncHandleRegistry.
-  bool registered_ = false;
+  SyncHandleRegistry::EventCallbackSubscription subscription_;
 
   // If non-zero, |event_| should be registered with SyncHandleRegistry.
   size_t register_request_count_ = 0;

mojo/public/cpp/bindings/sync_handle_registry.h

@@ -6,10 +6,12 @@
 #define MOJO_PUBLIC_CPP_BINDINGS_SYNC_HANDLE_REGISTRY_H_
 
 #include <map>
+#include <memory>
 
 #include "base/callback.h"
+#include "base/callback_helpers.h"
+#include "base/callback_list.h"
 #include "base/component_export.h"
-#include "base/containers/stack_container.h"
 #include "base/macros.h"
 #include "base/memory/ref_counted.h"
 #include "base/sequence_checker.h"
@@ -26,11 +28,31 @@ namespace mojo {
 class COMPONENT_EXPORT(MOJO_CPP_BINDINGS) SyncHandleRegistry
     : public base::RefCounted<SyncHandleRegistry> {
  public:
-  // Returns a sequence-local object.
-  static scoped_refptr<SyncHandleRegistry> current();
+  using EventCallbackList = base::RepeatingClosureList;
+  using EventCallback = EventCallbackList::CallbackType;
+
+  // Wrapper class that runs a closure after a CallbackList subscription is
+  // destroyed.
+  class COMPONENT_EXPORT(MOJO_CPP_BINDINGS) Subscription {
+   public:
+    Subscription(base::OnceClosure remove_closure,
+                 EventCallbackList* callbacks,
+                 EventCallback event_callback);
+    Subscription(Subscription&&);
+    Subscription& operator=(Subscription&&);
+    ~Subscription();
+
+   private:
+    base::ScopedClosureRunner remove_runner_;
+    std::unique_ptr<EventCallbackList::Subscription> subscription_;
+  };
+  using EventCallbackSubscription = std::unique_ptr<Subscription>;
 
   using HandleCallback = base::RepeatingCallback<void(MojoResult)>;
 
+  // Returns a sequence-local object.
+  static scoped_refptr<SyncHandleRegistry> current();
+
   // Registers a |Handle| to be watched for |handle_signals|. If any such
   // signals are satisfied during a Wait(), the Wait() is woken up and
   // |callback| is run.
@@ -44,12 +66,8 @@ class COMPONENT_EXPORT(MOJO_CPP_BINDINGS) SyncHandleRegistry
   // Wait() before any handle signals. |event| is not owned, and if it signals
   // during Wait(), |callback| is invoked.  Note that |event| may be registered
   // multiple times with different callbacks.
-  void RegisterEvent(base::WaitableEvent* event,
-                     base::RepeatingClosure callback);
-
-  // Unregisters a specific |event|+|callback| pair.
-  void UnregisterEvent(base::WaitableEvent* event,
-                       base::RepeatingClosure callback);
+  EventCallbackSubscription RegisterEvent(base::WaitableEvent* event,
+                                          EventCallback callback);
 
   // Waits on all the registered handles and events and runs callbacks
   // synchronously for any that become ready.
@@ -61,26 +79,17 @@ class COMPONENT_EXPORT(MOJO_CPP_BINDINGS) SyncHandleRegistry
  private:
   friend class base::RefCounted<SyncHandleRegistry>;
 
-  using EventCallbackList = base::StackVector<base::RepeatingClosure, 1>;
-  using EventMap = std::map<base::WaitableEvent*, EventCallbackList>;
-
   SyncHandleRegistry();
   ~SyncHandleRegistry();
 
-  void RemoveInvalidEventCallbacks();
-
   WaitSet wait_set_;
   std::map<Handle, HandleCallback> handles_;
-  EventMap events_;
-
-  // |true| iff this registry is currently dispatching event callbacks in
-  // Wait(). Used to allow for safe event registration/unregistration from event
-  // callbacks.
-  bool is_dispatching_event_callbacks_ = false;
+  std::map<base::WaitableEvent*, std::unique_ptr<EventCallbackList>> events_;
 
-  // Indicates if one or more event callbacks was unregistered during the most
-  // recent event callback dispatch.
-  bool remove_invalid_event_callbacks_after_dispatch_ = false;
+  // True when the registry is dispatching event callbacks in Wait().  This is
+  // used to improve the safety and efficiency of pruning unused entries in
+  // |events_| if Wait() results in reentrancy.
+  bool in_nested_wait_ = false;
 
   SEQUENCE_CHECKER(sequence_checker_);