Add waitForGracefulShutdown

Author: gjcairo

Sources/ServiceLifecycle/GracefulShutdown.swift

@@ -55,6 +55,18 @@ public func withGracefulShutdownHandler<T>(
     return try await operation()
 }
 
+/// Suspend the current task until graceful shutdown is executed.
+///
+/// A common use-case is to call this method in the `operation` closure of ``withGracefulShutdownHandler(operation:onGracefulShutdown:)``
+/// when you are only interested in executing a closure upon graceful shutdown, and nothing else.
+public func waitForGracefulShutdown() async {
+    guard let gracefulShutdownManager = TaskLocals.gracefulShutdownManager else {
+        return
+    }
+    
+    await gracefulShutdownManager.waitForGracefulShutdown()
+}
+
 /// This is just a helper type for the result of our task group.
 enum ValueOrGracefulShutdown<T: Sendable>: Sendable {
     case value(T)
@@ -138,6 +150,8 @@ public final class GracefulShutdownManager: @unchecked Sendable {
         fileprivate var handlerCounter: UInt64 = 0
         /// A boolean indicating if we have been shutdown already.
         fileprivate var isShuttingDown = false
+        /// Continuations to resume after all of the handlers have been executed.
+        fileprivate var gracefulShutdownFinishedContinuations = [CheckedContinuation<Void, Never>]()
     }
 
     private let state = LockedValueBox(State())
@@ -191,6 +205,25 @@ public final class GracefulShutdownManager: @unchecked Sendable {
             }
 
             state.handlers.removeAll()
+            
+            for continuation in state.gracefulShutdownFinishedContinuations {
+                continuation.resume()
+            }
+            
+            state.gracefulShutdownFinishedContinuations.removeAll()
+        }
+    }
+    
+    func waitForGracefulShutdown() async {
+        await withCheckedContinuation { continuation in
+            self.state.withLockedValue { state in
+                guard !state.isShuttingDown else {
+                    continuation.resume()
+                    return
+                }
+                
+                state.gracefulShutdownFinishedContinuations.append(continuation)
+            }
         }
     }
 }

Tests/ServiceLifecycleTests/GracefulShutdownTests.swift

@@ -253,4 +253,23 @@ final class GracefulShutdownTests: XCTestCase {
             XCTAssertTrue(Task.isShuttingDownGracefully)
         }
     }
+    
+    func testWaitForGracefulShutdown() async throws {
+        try await testGracefulShutdown { gracefulShutdownTestTrigger in
+            try await withThrowingTaskGroup(of: Void.self) { group in
+                group.addTask {
+                    try await Task.sleep(for: .milliseconds(10))
+                    gracefulShutdownTestTrigger.triggerGracefulShutdown()
+                }
+
+                await withGracefulShutdownHandler {
+                    await waitForGracefulShutdown()
+                } onGracefulShutdown: {
+                    // No-op
+                }
+
+                try await group.waitForAll()
+            }
+        }
+    }
 }