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,29 +667,28 @@ 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;
-    if (pump_messages_event) {
-      on_pump_messages_callback =
-          base::BindRepeating(&OnEventReady, &should_pump_messages);
-      registry->RegisterEvent(pump_messages_event, on_pump_messages_callback);
+      mojo::SyncHandleRegistry::EventCallbackSubscription
+          pump_messages_subsciption;
+      if (pump_messages_event) {
+        pump_messages_subsciption = registry->RegisterEvent(
+            pump_messages_event,
+            base::BindRepeating(&OnEventReady, &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();
     }
 
-    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.
       // We must process it to avoid potential deadlocks.

ipc/ipc_sync_message_filter.cc

@@ -69,25 +69,26 @@ 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);
+  {
+    bool done = false;
+    bool shutdown = false;
+    scoped_refptr<mojo::SyncHandleRegistry> registry =
+        mojo::SyncHandleRegistry::current();
+    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);
+    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_);
     delete pending_message.deserializer;

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;
+  // 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);
     }
-    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);
-    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_;
+  std::map<base::WaitableEvent*, std::unique_ptr<EventCallbackList>> 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;
-
-  // 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_);
 

mojo/public/cpp/bindings/tests/sync_handle_registry_unittest.cc

@@ -28,20 +28,20 @@ TEST_F(SyncHandleRegistryTest, DuplicateEventRegistration) {
   bool called1 = false;
   bool called2 = false;
   auto callback = [](bool* called) { *called = true; };
-  auto callback1 = base::BindRepeating(callback, &called1);
-  auto callback2 = base::BindRepeating(callback, &called2);
 
   base::WaitableEvent e(base::WaitableEvent::ResetPolicy::MANUAL,
                         base::WaitableEvent::InitialState::SIGNALED);
-  registry()->RegisterEvent(&e, callback1);
-  registry()->RegisterEvent(&e, callback2);
+  SyncHandleRegistry::EventCallbackSubscription subscription1 =
+      registry()->RegisterEvent(&e, base::BindRepeating(callback, &called1));
+  SyncHandleRegistry::EventCallbackSubscription subscription2 =
+      registry()->RegisterEvent(&e, base::BindRepeating(callback, &called2));
 
   const bool* stop_flags[] = {&called1, &called2};
   registry()->Wait(stop_flags, 2);
 
   EXPECT_TRUE(called1);
   EXPECT_TRUE(called2);
-  registry()->UnregisterEvent(&e, callback1);
+  subscription1.reset();
 
   called1 = false;
   called2 = false;
@@ -50,8 +50,6 @@ TEST_F(SyncHandleRegistryTest, DuplicateEventRegistration) {
 
   EXPECT_FALSE(called1);
   EXPECT_TRUE(called2);
-
-  registry()->UnregisterEvent(&e, callback2);
 }
 
 TEST_F(SyncHandleRegistryTest, UnregisterDuplicateEventInNestedWait) {
@@ -60,21 +58,25 @@ TEST_F(SyncHandleRegistryTest, UnregisterDuplicateEventInNestedWait) {
   bool called1 = false;
   bool called2 = false;
   bool called3 = false;
-  auto callback1 =
-      base::BindRepeating([](bool* called) { *called = true; }, &called1);
-  auto callback2 = base::BindRepeating(
-      [](base::WaitableEvent* e, base::RepeatingClosure other_callback,
-         scoped_refptr<SyncHandleRegistry> registry, bool* called) {
-        registry->UnregisterEvent(e, other_callback);
-        *called = true;
-      },
-      &e, callback1, registry(), &called2);
-  auto callback3 =
-      base::BindRepeating([](bool* called) { *called = true; }, &called3);
 
-  registry()->RegisterEvent(&e, callback1);
-  registry()->RegisterEvent(&e, callback2);
-  registry()->RegisterEvent(&e, callback3);
+  SyncHandleRegistry::EventCallbackSubscription subscription1 =
+      registry()->RegisterEvent(
+          &e,
+          base::BindRepeating([](bool* called) { *called = true; }, &called1));
+  SyncHandleRegistry::EventCallbackSubscription subscription2 =
+      registry()->RegisterEvent(
+          &e,
+          base::BindRepeating(
+              [](SyncHandleRegistry::EventCallbackSubscription* subscription,
+                 bool* called) {
+                subscription->reset();
+                *called = true;
+              },
+              &subscription1, &called2));
+  SyncHandleRegistry::EventCallbackSubscription subscription3 =
+      registry()->RegisterEvent(
+          &e,
+          base::BindRepeating([](bool* called) { *called = true; }, &called3));
 
   const bool* stop_flags[] = {&called1, &called2, &called3};
   registry()->Wait(stop_flags, 3);
@@ -90,7 +92,7 @@ TEST_F(SyncHandleRegistryTest, UnregisterDuplicateEventInNestedWait) {
   called2 = false;
   called3 = false;
 
-  registry()->UnregisterEvent(&e, callback2);
+  subscription2.reset();
   registry()->Wait(stop_flags, 3);
 
   EXPECT_FALSE(called1);
@@ -103,14 +105,14 @@ TEST_F(SyncHandleRegistryTest, UnregisterAndRegisterForNewEventInCallback) {
       base::WaitableEvent::ResetPolicy::MANUAL,
       base::WaitableEvent::InitialState::SIGNALED);
   bool called = false;
-  base::RepeatingClosure callback_holder;
+  SyncHandleRegistry::EventCallbackSubscription subscription;
   auto callback = base::BindRepeating(
       [](std::unique_ptr<base::WaitableEvent>* e,
-         base::RepeatingClosure* callback_holder,
+         SyncHandleRegistry::EventCallbackSubscription* subscription,
          scoped_refptr<SyncHandleRegistry> registry, bool* called) {
         EXPECT_FALSE(*called);
 
-        registry->UnregisterEvent(e->get(), *callback_holder);
+        subscription->reset();
         e->reset();
         *called = true;
 
@@ -118,17 +120,17 @@ TEST_F(SyncHandleRegistryTest, UnregisterAndRegisterForNewEventInCallback) {
             base::WaitableEvent::ResetPolicy::MANUAL,
             base::WaitableEvent::InitialState::SIGNALED);
         bool nested_called = false;
-        auto nested_callback = base::BindRepeating(
-            [](bool* called) { *called = true; }, &nested_called);
-        registry->RegisterEvent(&nested_event, nested_callback);
+        SyncHandleRegistry::EventCallbackSubscription nested_subscription =
+            registry->RegisterEvent(
+                &nested_event,
+                base::BindRepeating([](bool* called) { *called = true; },
+                                    &nested_called));
         const bool* stop_flag = &nested_called;
         registry->Wait(&stop_flag, 1);
-        registry->UnregisterEvent(&nested_event, nested_callback);
       },
-      &e, &callback_holder, registry(), &called);
-  callback_holder = callback;
+      &e, &subscription, registry(), &called);
 
-  registry()->RegisterEvent(e.get(), callback);
+  subscription = registry()->RegisterEvent(e.get(), callback);
 
   const bool* stop_flag = &called;
   registry()->Wait(&stop_flag, 1);
@@ -139,29 +141,29 @@ TEST_F(SyncHandleRegistryTest, UnregisterAndRegisterForSameEventInCallback) {
   base::WaitableEvent e(base::WaitableEvent::ResetPolicy::MANUAL,
                         base::WaitableEvent::InitialState::SIGNALED);
   bool called = false;
-  base::RepeatingClosure callback_holder;
+  SyncHandleRegistry::EventCallbackSubscription subscription;
   auto callback = base::BindRepeating(
-      [](base::WaitableEvent* e, base::RepeatingClosure* callback_holder,
+      [](base::WaitableEvent* e,
+         SyncHandleRegistry::EventCallbackSubscription* subscription,
          scoped_refptr<SyncHandleRegistry> registry, bool* called) {
         EXPECT_FALSE(*called);
 
-        registry->UnregisterEvent(e, *callback_holder);
+        subscription->reset();
         *called = true;
 
         bool nested_called = false;
-        auto nested_callback = base::BindRepeating(
-            [](bool* called) { *called = true; }, &nested_called);
-        registry->RegisterEvent(e, nested_callback);
+        SyncHandleRegistry::EventCallbackSubscription nested_subscription =
+            registry->RegisterEvent(
+                e, base::BindRepeating([](bool* called) { *called = true; },
+                                       &nested_called));
         const bool* stop_flag = &nested_called;
         registry->Wait(&stop_flag, 1);
-        registry->UnregisterEvent(e, nested_callback);
 
         EXPECT_TRUE(nested_called);
       },
-      &e, &callback_holder, registry(), &called);
-  callback_holder = callback;
+      &e, &subscription, registry(), &called);
 
-  registry()->RegisterEvent(&e, callback);
+  subscription = registry()->RegisterEvent(&e, callback);
 
   const bool* stop_flag = &called;
   registry()->Wait(&stop_flag, 1);
@@ -184,26 +186,24 @@ TEST_F(SyncHandleRegistryTest, RegisterDuplicateEventFromWithinCallback) {
         *called = true;
 
         bool called2 = false;
-        auto callback2 =
-            base::BindRepeating([](bool* called) { *called = true; }, &called2);
-        registry->RegisterEvent(e, callback2);
+        SyncHandleRegistry::EventCallbackSubscription nested_subscription =
+            registry->RegisterEvent(
+                e, base::BindRepeating([](bool* called) { *called = true; },
+                                       &called2));
 
         const bool* stop_flag = &called2;
         registry->Wait(&stop_flag, 1);
-
-        registry->UnregisterEvent(e, callback2);
       },
       &e, registry(), &called, &call_count);
 
-  registry()->RegisterEvent(&e, callback);
+  SyncHandleRegistry::EventCallbackSubscription subscription =
+      registry()->RegisterEvent(&e, callback);
 
   const bool* stop_flag = &called;
   registry()->Wait(&stop_flag, 1);
 
   EXPECT_TRUE(called);
   EXPECT_EQ(2, call_count);
-
-  registry()->UnregisterEvent(&e, callback);
 }
 
 TEST_F(SyncHandleRegistryTest, UnregisterUniqueEventInNestedWait) {
@@ -214,27 +214,31 @@ TEST_F(SyncHandleRegistryTest, UnregisterUniqueEventInNestedWait) {
                          base::WaitableEvent::InitialState::SIGNALED);
   bool called1 = false;
   bool called2 = false;
-  auto callback1 =
-      base::BindRepeating([](bool* called) { *called = true; }, &called1);
+
+  SyncHandleRegistry::EventCallbackSubscription subscription1 =
+      registry()->RegisterEvent(
+          e1.get(),
+          base::BindRepeating([](bool* called) { *called = true; }, &called1));
   auto callback2 = base::BindRepeating(
       [](std::unique_ptr<base::WaitableEvent>* e1,
-         base::RepeatingClosure other_callback,
+         SyncHandleRegistry::EventCallbackSubscription* subscription,
          scoped_refptr<SyncHandleRegistry> registry, bool* called) {
         // Prevent re-entrancy.
         if (*called)
           return;
 
-        registry->UnregisterEvent(e1->get(), other_callback);
+        subscription->reset();
         *called = true;
         e1->reset();
 
         // Nest another wait.
         bool called3 = false;
-        auto callback3 =
-            base::BindRepeating([](bool* called) { *called = true; }, &called3);
         base::WaitableEvent e3(base::WaitableEvent::ResetPolicy::MANUAL,
                                base::WaitableEvent::InitialState::SIGNALED);
-        registry->RegisterEvent(&e3, callback3);
+        SyncHandleRegistry::EventCallbackSubscription nested_subscription =
+            registry->RegisterEvent(
+                &e3, base::BindRepeating([](bool* called) { *called = true; },
+                                         &called3));
 
         // This nested Wait() must not attempt to wait on |e1| since it has
         // been unregistered. This would crash otherwise, since |e1| has been
@@ -243,19 +247,16 @@ TEST_F(SyncHandleRegistryTest, UnregisterUniqueEventInNestedWait) {
         registry->Wait(stop_flags, 1);
 
         EXPECT_TRUE(called3);
-        registry->UnregisterEvent(&e3, callback3);
       },
-      &e1, callback1, registry(), &called2);
+      &e1, &subscription1, registry(), &called2);
 
-  registry()->RegisterEvent(e1.get(), callback1);
-  registry()->RegisterEvent(&e2, callback2);
+  SyncHandleRegistry::EventCallbackSubscription subscription2 =
+      registry()->RegisterEvent(&e2, callback2);
 
   const bool* stop_flags[] = {&called1, &called2};
   registry()->Wait(stop_flags, 2);
 
   EXPECT_TRUE(called2);
-
-  registry()->UnregisterEvent(&e2, callback2);
 }
 
 }  // namespace mojo