From 8bd5d484fe558bf4902e34f66ef138c6e1581c8b Mon Sep 17 00:00:00 2001
From: Tomas Gunnarsson <tommi@webrtc.org>
Date: Sat, 9 Apr 2022 17:49:00 +0200
Subject: [PATCH] [CallbackList] Allow removal from within callback.

This makes moving from sigslot to CallbackList slightly simpler in some
situations.

Bug: webrtc:11943
Change-Id: I5c6dafb8ac597a119b90b64f369fa9e6316e38da
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/258560
Reviewed-by: Mirko Bonadei <mbonadei@webrtc.org>
Commit-Queue: Tomas Gunnarsson <tommi@webrtc.org>
Cr-Commit-Position: refs/heads/main@{#36515}
---
 rtc_base/callback_list.cc          | 24 +++++++++++++++++++-----
 rtc_base/callback_list.h           |  8 ++++++++
 rtc_base/callback_list_unittest.cc | 13 +++++++++++++
 3 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/rtc_base/callback_list.cc b/rtc_base/callback_list.cc
index 88d0b6fc71..c452c79b38 100644
--- a/rtc_base/callback_list.cc
+++ b/rtc_base/callback_list.cc
@@ -22,8 +22,7 @@ CallbackListReceivers::~CallbackListReceivers() {
 }
 
 void CallbackListReceivers::RemoveReceivers(const void* removal_tag) {
-  RTC_CHECK(!send_in_progress_);
-  RTC_DCHECK(removal_tag != nullptr);
+  RTC_DCHECK(removal_tag);
 
   // We divide the receivers_ vector into three regions: from right to left, the
   // "keep" region, the "todo" region, and the "remove" region. The "todo"
@@ -42,8 +41,13 @@ void CallbackListReceivers::RemoveReceivers(const void* removal_tag) {
     } else if (receivers_[first_remove - 1].removal_tag == removal_tag) {
       // The last element of the "todo" region should be removed. Move the
       // "todo"/"remove" boundary.
+      if (send_in_progress_) {
+        // Tag this receiver for removal, which will be done when `ForEach`
+        // has completed.
+        receivers_[first_remove - 1].removal_tag = pending_removal_tag();
+      }
       --first_remove;
-    } else {
+    } else if (!send_in_progress_) {
       // The first element of the "todo" region should be removed, and the last
       // element of the "todo" region should be kept. Swap them, and then shrink
       // the "todo" region from both ends.
@@ -57,18 +61,28 @@ void CallbackListReceivers::RemoveReceivers(const void* removal_tag) {
     }
   }
 
-  // Discard the remove region.
-  receivers_.resize(first_remove);
+  if (!send_in_progress_) {
+    // Discard the remove region.
+    receivers_.resize(first_remove);
+  }
 }
 
 void CallbackListReceivers::Foreach(
     rtc::FunctionView<void(UntypedFunction&)> fv) {
   RTC_CHECK(!send_in_progress_);
+  bool removals_detected = false;
   send_in_progress_ = true;
   for (auto& r : receivers_) {
+    RTC_DCHECK_NE(r.removal_tag, pending_removal_tag());
     fv(r.function);
+    if (r.removal_tag == pending_removal_tag()) {
+      removals_detected = true;
+    }
   }
   send_in_progress_ = false;
+  if (removals_detected) {
+    RemoveReceivers(pending_removal_tag());
+  }
 }
 
 template void CallbackListReceivers::AddReceiver(
diff --git a/rtc_base/callback_list.h b/rtc_base/callback_list.h
index 18d48b02ee..32ea1b85f9 100644
--- a/rtc_base/callback_list.h
+++ b/rtc_base/callback_list.h
@@ -51,10 +51,18 @@ class CallbackListReceivers {
   void Foreach(rtc::FunctionView<void(UntypedFunction&)> fv);
 
  private:
+  // Special protected pointer value that's used as a removal_tag for
+  // receivers that want to unsubscribe from within a callback.
+  // Note we could use `&receivers_` too, but since it's the first member
+  // variable of the class, its address will be the same as the instance
+  // CallbackList instance, so we take an extra step to avoid collision.
+  const void* pending_removal_tag() const { return &send_in_progress_; }
+
   struct Callback {
     const void* removal_tag;
     UntypedFunction function;
   };
+
   std::vector<Callback> receivers_;
   bool send_in_progress_ = false;
 };
diff --git a/rtc_base/callback_list_unittest.cc b/rtc_base/callback_list_unittest.cc
index e2bc6d515e..483eb3f99a 100644
--- a/rtc_base/callback_list_unittest.cc
+++ b/rtc_base/callback_list_unittest.cc
@@ -252,5 +252,18 @@ TEST(CallbackList, RemoveManyReceivers) {
   EXPECT_EQ(accumulator, 1212);
 }
 
+TEST(CallbackList, RemoveFromSend) {
+  int removal_tag = 0;
+  CallbackList<> c;
+  c.AddReceiver(&removal_tag, [&] {
+    c.RemoveReceivers(&removal_tag);
+    // Do after RemoveReceivers to make sure the lambda is still valid.
+    ++removal_tag;
+  });
+  c.Send();
+  c.Send();
+  EXPECT_EQ(removal_tag, 1);
+}
+
 }  // namespace
 }  // namespace webrtc