From ac9e52181bbdec1df41e99dce2eb748843719227 Mon Sep 17 00:00:00 2001
From: Craig Tiller <ctiller@google.com>
Date: Mon, 13 Sep 2021 17:39:02 -0700
Subject: [PATCH] Activities: Add force wakeup API (#27336)

* Activities: Add force wakeup API

* Automated change: Fix sanity tests

Co-authored-by: ctiller <ctiller@users.noreply.github.com>
---
 src/core/lib/promise/activity.h    |  7 +++++++
 test/core/promise/activity_test.cc | 21 +++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/core/lib/promise/activity.h b/src/core/lib/promise/activity.h
index e9a1b73fb9e..76ef7f7f9fd 100644
--- a/src/core/lib/promise/activity.h
+++ b/src/core/lib/promise/activity.h
@@ -113,6 +113,13 @@ class Activity : private Wakeable {
   // Fetch the size of the implementation of this activity.
   virtual size_t Size() = 0;
 
+  // Force wakeup from the outside.
+  // This should be rarely needed, and usages should be accompanied with a note
+  // on why it's not possible to wakeup with a Waker object.
+  // Nevertheless, it's sometimes useful for integrations with Activity to force
+  // an Activity to repoll.
+  void ForceWakeup() { MakeOwningWaker().Wakeup(); }
+
   // Wakeup the current threads activity - will force a subsequent poll after
   // the one that's running.
   static void WakeupCurrent() { current()->got_wakeup_during_run_ = true; }
diff --git a/test/core/promise/activity_test.cc b/test/core/promise/activity_test.cc
index 854f6549ac1..3d9c9a1d968 100644
--- a/test/core/promise/activity_test.cc
+++ b/test/core/promise/activity_test.cc
@@ -231,6 +231,27 @@ TYPED_TEST(BarrierTest, WakeAfterDestruction) {
   b.Clear();
 }
 
+TEST(ActivityTest, ForceWakeup) {
+  StrictMock<MockFunction<void(absl::Status)>> on_done;
+  int run_count = 0;
+  auto activity = MakeActivity(
+      [&run_count]() -> Poll<absl::Status> {
+        ++run_count;
+        switch (run_count) {
+          case 1:
+            return Pending{};
+          case 2:
+            return absl::OkStatus();
+          default:
+            abort();
+        }
+      },
+      NoCallbackScheduler(),
+      [&on_done](absl::Status status) { on_done.Call(std::move(status)); });
+  EXPECT_CALL(on_done, Call(absl::OkStatus()));
+  activity->ForceWakeup();
+}
+
 struct TestContext {
   bool* done;
 };