fix multiple wakeup bug

ctiller · Dec 2, 2024 · 42c9cbf · 42c9cbf
1 parent 721ab4f
commit 42c9cbf
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/src/core/lib/promise/party.h b/src/core/lib/promise/party.h
@@ -95,7 +95,8 @@ class Party : public Activity, private Wakeable {
     explicit WakeupHold(Party* party)
         : prev_state_(party->state_.load(std::memory_order_relaxed)) {
       // Try to lock
-      if (party->state_.compare_exchange_weak(prev_state_,
+      if ((prev_state_ & kLocked) == 0 &&
+          party->state_.compare_exchange_weak(prev_state_,
                                               (prev_state_ | kLocked) + kOneRef,
                                               std::memory_order_relaxed)) {
         DCHECK_EQ(prev_state_ & ~(kRefMask | kAllocatedMask), 0u)

diff --git a/test/core/promise/party_test.cc b/test/core/promise/party_test.cc
@@ -268,6 +268,26 @@ TEST_F(PartyTest, CanBulkSpawn) {
   n2.WaitForNotification();
 }
 
+TEST_F(PartyTest, CanNestWakeupHold) {
+  auto party = MakeParty();
+  Notification n1;
+  Notification n2;
+  {
+    Party::WakeupHold hold1(party.get());
+    Party::WakeupHold hold2(party.get());
+    party->Spawn(
+        "spawn1", []() { return Empty{}; }, [&n1](Empty) { n1.Notify(); });
+    party->Spawn(
+        "spawn2", []() { return Empty{}; }, [&n2](Empty) { n2.Notify(); });
+    for (int i = 0; i < 5000; i++) {
+      EXPECT_FALSE(n1.HasBeenNotified());
+      EXPECT_FALSE(n2.HasBeenNotified());
+    }
+  }
+  n1.WaitForNotification();
+  n2.WaitForNotification();
+}
+
 TEST_F(PartyTest, ThreadStressTest) {
   auto party = MakeParty();
   std::vector<std::thread> threads;