From ab5259b290c7a298cde87c5252b174d968d1cd7c Mon Sep 17 00:00:00 2001
From: Martin Chang <marty1885@users.noreply.github.com>
Date: Thu, 9 Mar 2023 10:13:18 +0800
Subject: [PATCH] Add coroutine to wait until event loop ends (#1500)

---
 lib/inc/drogon/utils/coroutine.h | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/lib/inc/drogon/utils/coroutine.h b/lib/inc/drogon/utils/coroutine.h
index 6a35134f..09148897 100644
--- a/lib/inc/drogon/utils/coroutine.h
+++ b/lib/inc/drogon/utils/coroutine.h
@@ -697,6 +697,21 @@ struct [[nodiscard]] LoopAwaiter : CallbackAwaiter<void>
     std::function<void()> taskFunc_;
 };
 
+struct [[nodiscard]] EndAwaiter : CallbackAwaiter<void>
+{
+    EndAwaiter(trantor::EventLoop *loop) : loop_(loop)
+    {
+        assert(loop);
+    }
+    void await_suspend(std::coroutine_handle<> handle)
+    {
+        loop_->runOnQuit([handle]() { handle.resume(); });
+    }
+
+  private:
+    trantor::EventLoop *loop_{nullptr};
+};
+
 }  // namespace internal
 
 inline internal::TimerAwaiter sleepCoro(
@@ -723,6 +738,12 @@ inline internal::LoopAwaiter queueInLoopCoro(
     return {workLoop, std::move(taskFunc), resumeLoop};
 }
 
+inline internal::EndAwaiter untilQuit(trantor::EventLoop *loop)
+{
+    assert(loop);
+    return {loop};
+}
+
 template <typename T, typename = std::void_t<>>
 struct is_resumable : std::false_type
 {