From 8da5953fb28232235ceda85d66f0f847c3964176 Mon Sep 17 00:00:00 2001
From: Tommi <tommi@webrtc.org>
Date: Tue, 24 Oct 2023 23:43:56 +0200
Subject: [PATCH] Support initializing PendingTaskSafetyFlag with a specific
 TaskQueue.

Bug: none
Change-Id: I0f354708e6275372601adc36da3012259bb57303
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/324280
Reviewed-by: Danil Chapovalov <danilchap@webrtc.org>
Commit-Queue: Tomas Gunnarsson <tommi@webrtc.org>
Cr-Commit-Position: refs/heads/main@{#41008}
---
 api/task_queue/pending_task_safety_flag.cc          | 10 ++++++++++
 api/task_queue/pending_task_safety_flag.h           |  8 ++++++++
 api/task_queue/pending_task_safety_flag_unittest.cc | 11 +++++++++++
 3 files changed, 29 insertions(+)

diff --git a/api/task_queue/pending_task_safety_flag.cc b/api/task_queue/pending_task_safety_flag.cc
index 437ce0755d..3145550829 100644
--- a/api/task_queue/pending_task_safety_flag.cc
+++ b/api/task_queue/pending_task_safety_flag.cc
@@ -32,6 +32,16 @@ PendingTaskSafetyFlag::CreateDetached() {
   return safety_flag;
 }
 
+// Creates a flag, but with its SequenceChecker explicitly initialized for
+// a given task queue and the `alive()` flag specified.
+rtc::scoped_refptr<PendingTaskSafetyFlag>
+PendingTaskSafetyFlag::CreateAttachedToTaskQueue(
+    bool alive,
+    TaskQueueBase* attached_queue) {
+  return rtc::scoped_refptr<PendingTaskSafetyFlag>(
+      new PendingTaskSafetyFlag(alive, attached_queue));
+}
+
 rtc::scoped_refptr<PendingTaskSafetyFlag>
 PendingTaskSafetyFlag::CreateDetachedInactive() {
   rtc::scoped_refptr<PendingTaskSafetyFlag> safety_flag = CreateInternal(false);
diff --git a/api/task_queue/pending_task_safety_flag.h b/api/task_queue/pending_task_safety_flag.h
index 7f6a592856..b601b94215 100644
--- a/api/task_queue/pending_task_safety_flag.h
+++ b/api/task_queue/pending_task_safety_flag.h
@@ -68,6 +68,12 @@ class RTC_EXPORT PendingTaskSafetyFlag final
   // may be created on a different thread than the flag will be used on.
   static rtc::scoped_refptr<PendingTaskSafetyFlag> CreateDetached();
 
+  // Creates a flag, but with its SequenceChecker explicitly initialized for
+  // a given task queue and the `alive()` flag specified.
+  static rtc::scoped_refptr<PendingTaskSafetyFlag> CreateAttachedToTaskQueue(
+      bool alive,
+      TaskQueueBase* attached_queue);
+
   // Same as `CreateDetached()` except the initial state of the returned flag
   // will be `!alive()`.
   static rtc::scoped_refptr<PendingTaskSafetyFlag> CreateDetachedInactive();
@@ -95,6 +101,8 @@ class RTC_EXPORT PendingTaskSafetyFlag final
 
  protected:
   explicit PendingTaskSafetyFlag(bool alive) : alive_(alive) {}
+  PendingTaskSafetyFlag(bool alive, TaskQueueBase* attached_queue)
+      : alive_(alive), main_sequence_(attached_queue) {}
 
  private:
   static rtc::scoped_refptr<PendingTaskSafetyFlag> CreateInternal(bool alive);
diff --git a/api/task_queue/pending_task_safety_flag_unittest.cc b/api/task_queue/pending_task_safety_flag_unittest.cc
index cedf0eb8df..3a1ed268fc 100644
--- a/api/task_queue/pending_task_safety_flag_unittest.cc
+++ b/api/task_queue/pending_task_safety_flag_unittest.cc
@@ -167,6 +167,17 @@ TEST(PendingTaskSafetyFlagTest, PendingTaskNotAliveInitialized) {
   EXPECT_TRUE(task_2_ran);
 }
 
+TEST(PendingTaskSafetyFlagTest, PendingTaskInitializedForTaskQueue) {
+  TaskQueueForTest tq("PendingTaskAliveInitializedForTaskQueue");
+
+  // Create a new flag that initially `alive`, attached to a specific TQ.
+  auto flag = PendingTaskSafetyFlag::CreateAttachedToTaskQueue(true, tq.Get());
+  tq.SendTask([&flag]() { EXPECT_TRUE(flag->alive()); });
+  // Repeat the same steps but initialize as inactive.
+  flag = PendingTaskSafetyFlag::CreateAttachedToTaskQueue(false, tq.Get());
+  tq.SendTask([&flag]() { EXPECT_FALSE(flag->alive()); });
+}
+
 TEST(PendingTaskSafetyFlagTest, SafeTask) {
   rtc::scoped_refptr<PendingTaskSafetyFlag> flag =
       PendingTaskSafetyFlag::Create();