Add yielding support to ScopedOperationsBatcher

Also including a change for the internal BatchedTask struct and use
std::variant instead as suggested in a previous CL.

Bug: webrtc:42222804
Change-Id: Ia8d10cb9ee6c0a735420232aeca82fb6d85233ca
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/454940
Reviewed-by: Harald Alvestrand <hta@webrtc.org>
Commit-Queue: Tomas Gunnarsson <tommi@webrtc.org>
Cr-Commit-Position: refs/heads/main@{#47201}
diff --git a/pc/BUILD.gn b/pc/BUILD.gn
index 2755d73..88287d5 100644
--- a/pc/BUILD.gn
+++ b/pc/BUILD.gn
@@ -1076,8 +1076,10 @@
     "scoped_operations_batcher.h",
   ]
   deps = [
+    "../api:sequence_checker",
     "../rtc_base:checks",
     "../rtc_base:threading",
+    "../rtc_base/system:no_unique_address",
     "//third_party/abseil-cpp/absl/functional:any_invocable",
   ]
 }
diff --git a/pc/scoped_operations_batcher.cc b/pc/scoped_operations_batcher.cc
index 050d0c3..00c19cf 100644
--- a/pc/scoped_operations_batcher.cc
+++ b/pc/scoped_operations_batcher.cc
@@ -10,18 +10,21 @@
 
 #include "pc/scoped_operations_batcher.h"
 
+#include <cstddef>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include "absl/functional/any_invocable.h"
+#include "api/sequence_checker.h"
 #include "rtc_base/checks.h"
 #include "rtc_base/thread.h"
 
 namespace webrtc {
 
-ScopedOperationsBatcher::ScopedOperationsBatcher(Thread* worker_thread)
-    : worker_thread_(worker_thread) {
-  RTC_DCHECK(worker_thread_);
+ScopedOperationsBatcher::ScopedOperationsBatcher(Thread* target_thread)
+    : target_thread_(target_thread) {
+  RTC_DCHECK(target_thread_);
 }
 
 ScopedOperationsBatcher::~ScopedOperationsBatcher() {
@@ -29,39 +32,56 @@
 }
 
 void ScopedOperationsBatcher::Run() {
-  std::vector<absl::AnyInvocable<void() &&>> signaling_tasks;
-  if (!tasks_.empty()) {
-    worker_thread_->BlockingCall([&] {
-      for (auto& task : tasks_) {
-        if (task.void_task) {
-          std::move(task.void_task)();
+  RTC_DCHECK_RUN_ON(&sequence_checker_);
+  std::vector<absl::AnyInvocable<void() &&>> return_tasks;
+
+  size_t task_idx = 0;
+  bool target_thread_is_current = target_thread_->IsCurrent();
+
+  while (task_idx < tasks_.size()) {
+    target_thread_->BlockingCall([&] {
+      while (task_idx < tasks_.size()) {
+        if (auto* void_task =
+                std::get_if<absl::AnyInvocable<void() &&>>(&tasks_[task_idx])) {
+          std::move (*void_task)();
         } else {
-          RTC_DCHECK(task.returning_task);
-          auto ret = std::move(task.returning_task)();
+          auto* returning_task = std::get_if<
+              absl::AnyInvocable<absl::AnyInvocable<void() &&>() &&>>(
+              &tasks_[task_idx]);
+          RTC_DCHECK(returning_task);
+          auto ret = std::move(*returning_task)();
           if (ret) {
-            signaling_tasks.push_back(std::move(ret));
+            return_tasks.push_back(std::move(ret));
           }
         }
+        ++task_idx;
+        if (!target_thread_is_current && target_thread_->IsYieldRequested()) {
+          return;
+        }
       }
     });
-    tasks_.clear();
   }
 
-  for (auto& task : signaling_tasks) {
+  RTC_DCHECK_EQ(task_idx, tasks_.size());
+  tasks_.clear();
+
+  for (auto& task : return_tasks) {
     std::move(task)();
   }
 }
 
 void ScopedOperationsBatcher::push_back(absl::AnyInvocable<void() &&> task) {
+  RTC_DCHECK_RUN_ON(&sequence_checker_);
   if (task) {
-    tasks_.push_back({.void_task = std::move(task)});
+    tasks_.emplace_back(std::move(task));
   }
 }
 
 void ScopedOperationsBatcher::push_back(
     absl::AnyInvocable<absl::AnyInvocable<void() &&>() &&> task) {
+  RTC_DCHECK_RUN_ON(&sequence_checker_);
   if (task) {
-    tasks_.push_back({.returning_task = std::move(task)});
+    tasks_.emplace_back(std::move(task));
   }
 }
 
diff --git a/pc/scoped_operations_batcher.h b/pc/scoped_operations_batcher.h
index d9a4c23..f5c32c1 100644
--- a/pc/scoped_operations_batcher.h
+++ b/pc/scoped_operations_batcher.h
@@ -11,17 +11,25 @@
 #ifndef PC_SCOPED_OPERATIONS_BATCHER_H_
 #define PC_SCOPED_OPERATIONS_BATCHER_H_
 
+#include <variant>
 #include <vector>
 
 #include "absl/functional/any_invocable.h"
+#include "api/sequence_checker.h"
+#include "rtc_base/system/no_unique_address.h"
 #include "rtc_base/thread.h"
 
 namespace webrtc {
 
-// Batches operations to be executed synchronously on the worker thread.
+// Batches operations to be executed synchronously on a target thread.
+//
+// ScopedOperationsBatcher must only be created on the signaling thread.
 //
 // When the batcher goes out of scope (or `Run()` is explicitly called), it
-// executes all queued tasks in a single `BlockingCall` to the worker thread.
+// executes all queued tasks in one or more `BlockingCall`s to the target
+// thread. If the target thread requests a yield during batch execution, the
+// batcher will cooperatively yield the thread and resume execution in a
+// subsequent `BlockingCall`.
 //
 // Tasks can either have a `void` return type, or return a new task.
 // Any tasks returned by the executed worker thread tasks are collected
@@ -29,23 +37,23 @@
 // thread) after the worker thread operations have completed.
 class ScopedOperationsBatcher {
  public:
-  explicit ScopedOperationsBatcher(Thread* worker_thread);
+  explicit ScopedOperationsBatcher(Thread* target_thread);
   ~ScopedOperationsBatcher();
 
   void Run();
 
-  // Queues non-nullptr tasks to be executed on the worker when the
+  // Queues non-nullptr tasks to be executed on the target thread when the
   // ScopedOperationsBatcher goes out of scope.
   void push_back(absl::AnyInvocable<void() &&> task);
   void push_back(absl::AnyInvocable<absl::AnyInvocable<void() &&>() &&> task);
 
  private:
-  struct BatchedTask {
-    absl::AnyInvocable<void() &&> void_task;
-    absl::AnyInvocable<absl::AnyInvocable<void() &&>() &&> returning_task;
-  };
+  using BatchedTask =
+      std::variant<absl::AnyInvocable<void() &&>,
+                   absl::AnyInvocable<absl::AnyInvocable<void() &&>() &&>>;
 
-  Thread* const worker_thread_;
+  RTC_NO_UNIQUE_ADDRESS SequenceChecker sequence_checker_;
+  Thread* const target_thread_;
   std::vector<BatchedTask> tasks_;
 };
 
diff --git a/pc/scoped_operations_batcher_unittest.cc b/pc/scoped_operations_batcher_unittest.cc
index 7f15ad3..0f188f7 100644
--- a/pc/scoped_operations_batcher_unittest.cc
+++ b/pc/scoped_operations_batcher_unittest.cc
@@ -12,6 +12,7 @@
 
 #include <memory>
 #include <utility>
+#include <vector>
 
 #include "absl/functional/any_invocable.h"
 #include "rtc_base/thread.h"
@@ -20,23 +21,23 @@
 namespace webrtc {
 namespace {
 
-TEST(ScopedOperationsBatcherTest, ExecutesTasksOnWorkerThread) {
-  auto worker_thread = Thread::Create();
-  worker_thread->Start();
+TEST(ScopedOperationsBatcherTest, ExecutesTasksOnTargetThread) {
+  auto target_thread = Thread::Create();
+  target_thread->Start();
 
   bool task_executed = false;
-  bool worker_checked = false;
+  bool target_checked = false;
 
   {
-    ScopedOperationsBatcher batcher(worker_thread.get());
+    ScopedOperationsBatcher batcher(target_thread.get());
     batcher.push_back([&] {
       task_executed = true;
-      worker_checked = worker_thread->IsCurrent();
+      target_checked = target_thread->IsCurrent();
     });
   }
 
   EXPECT_TRUE(task_executed);
-  EXPECT_TRUE(worker_checked);
+  EXPECT_TRUE(target_checked);
 }
 
 TEST(ScopedOperationsBatcherTest, ExecutesReturnedTasksOnCallingThread) {
@@ -46,8 +47,8 @@
   // `~ScopedOperationsBatcher()`.
   auto signaling_thread = Thread::Current();
 
-  auto worker_thread = Thread::Create();
-  worker_thread->Start();
+  auto target_thread = Thread::Create();
+  target_thread->Start();
 
   bool return_task_executed = false;
   Thread* return_task_thread = nullptr;
@@ -55,7 +56,7 @@
   Thread* task_thread = nullptr;
 
   {
-    ScopedOperationsBatcher batcher(worker_thread.get());
+    ScopedOperationsBatcher batcher(target_thread.get());
     absl::AnyInvocable<absl::AnyInvocable<void() &&>() &&> task =
         [&]() -> absl::AnyInvocable<void() &&> {
       task_executed = true;
@@ -69,10 +70,34 @@
   }
 
   EXPECT_TRUE(task_executed);
-  EXPECT_EQ(task_thread, worker_thread.get());
+  EXPECT_EQ(task_thread, target_thread.get());
   EXPECT_TRUE(return_task_executed);
   EXPECT_EQ(return_task_thread, signaling_thread);
 }
 
+TEST(ScopedOperationsBatcherTest, YieldsToHighPriorityTasks) {
+  auto target_thread = Thread::Create();
+  target_thread->Start();
+
+  std::vector<int> execution_order;
+
+  {
+    ScopedOperationsBatcher batcher(target_thread.get());
+    batcher.push_back([&] { execution_order.push_back(1); });
+    batcher.push_back([&] {
+      execution_order.push_back(2);
+      // Post a high-priority task that should interrupt the batch.
+      target_thread->PostHighPriorityTask(
+          [&] { execution_order.push_back(3); });
+    });
+    batcher.push_back([&] { execution_order.push_back(4); });
+    batcher.push_back([&] { execution_order.push_back(5); });
+  }
+
+  // Expect the high priority task (3) to execute immediately after the task
+  // that posted it (2). The regular tasks (4, 5) should yield.
+  EXPECT_EQ(execution_order, std::vector<int>({1, 2, 3, 4, 5}));
+}
+
 }  // namespace
 }  // namespace webrtc