Add extensions to TaskGroup and ThrowingTaskGroup

Author: FranzBusch

Sources/ServiceLifecycle/GracefulShutdown.swift

@@ -48,7 +48,6 @@ extension Task {
     /// Indicates that the task should shutdown gracefully.
     ///
     /// Calling this method on a task that doesn’t support graceful shutdown has no effect.
-    /// Likewise, if the task has already run past the last point where it would stop early, calling this method has no effect.
     ///
     /// - Note: This method is mostly relevant for testing graceful shutdown. In your application, the ``ServiceRunner``
     /// should be the instance that triggers the graceful shutdown.
@@ -62,6 +61,40 @@ extension Task {
     }
 }
 
+extension TaskGroup {
+    /// Indicates all of the tasks in the group to shutdown gracefully.
+    ///
+    /// Calling this method on a task that doesn’t support graceful shutdown has no effect.
+    ///
+    /// - Note: This method is mostly relevant for testing graceful shutdown. In your application, the ``ServiceRunner``
+    /// should be the instance that triggers the graceful shutdown.
+    public func shutdownGracefullyAll() async {
+        guard let gracefulShutdownManager = TaskLocals.gracefulShutdownManager else {
+            print("WARNING: Trying to shutdown gracefully inside a task that doesn't have access to the ShutdownGracefulManager. This happens either when unstructured Concurrency is used like Task.detached {} or when you tried to shutdown gracefully outside the ServiceRunner.run method.")
+            return
+        }
+
+        await gracefulShutdownManager.shutdownGracefully()
+    }
+}
+
+extension ThrowingTaskGroup {
+    /// Indicates all of the tasks in the group to shutdown gracefully.
+    ///
+    /// Calling this method on a task that doesn’t support graceful shutdown has no effect.
+    ///
+    /// - Note: This method is mostly relevant for testing graceful shutdown. In your application, the ``ServiceRunner``
+    /// should be the instance that triggers the graceful shutdown.
+    public func shutdownGracefullyAll() async {
+        guard let gracefulShutdownManager = TaskLocals.gracefulShutdownManager else {
+            print("WARNING: Trying to shutdown gracefully inside a task that doesn't have access to the ShutdownGracefulManager. This happens either when unstructured Concurrency is used like Task.detached {} or when you tried to shutdown gracefully outside the ServiceRunner.run method.")
+            return
+        }
+
+        await gracefulShutdownManager.shutdownGracefully()
+    }
+}
+
 @_spi(Testing)
 public actor GracefulShutdownManager {
     struct Handler {

Tests/ServiceLifecycleTests/GracefulShutdownTests.swift

@@ -134,4 +134,89 @@ final class GracefulShutdownTests: XCTestCase {
             XCTAssertNil(weakFoo)
         }
     }
+
+    func testTaskShutdownGracefully() async {
+        let shutdownGracefulManager = GracefulShutdownManager()
+        await TaskLocals.$gracefulShutdownManager.withValue(shutdownGracefulManager) {
+            let task = Task {
+                var cont: AsyncStream<Void>.Continuation!
+                let stream = AsyncStream<Void> { cont = $0 }
+                let continuation = cont!
+
+                await withShutdownGracefulHandler {
+                    await withTaskGroup(of: Void.self) { group in
+                        group.addTask {
+                            await stream.first { _ in true }
+                        }
+
+                        await group.waitForAll()
+                    }
+                } onGracefulShutdown: {
+                    continuation.finish()
+                }
+            }
+
+            await task.shutdownGracefully()
+
+            await task.value
+        }
+    }
+
+    func testTaskGroupShutdownGracefully() async {
+        let shutdownGracefulManager = GracefulShutdownManager()
+        await TaskLocals.$gracefulShutdownManager.withValue(shutdownGracefulManager) {
+            await withTaskGroup(of: Void.self) { group in
+                group.addTask {
+                    var cont: AsyncStream<Void>.Continuation!
+                    let stream = AsyncStream<Void> { cont = $0 }
+                    let continuation = cont!
+
+                    await withShutdownGracefulHandler {
+                        await withTaskGroup(of: Void.self) { group in
+                            group.addTask {
+                                await stream.first { _ in true }
+                            }
+
+                            await group.waitForAll()
+                        }
+                    } onGracefulShutdown: {
+                        continuation.finish()
+                    }
+                }
+
+                await group.shutdownGracefullyAll()
+
+                await group.waitForAll()
+            }
+        }
+    }
+
+    func testThrowingTaskGroupShutdownGracefully() async {
+        let shutdownGracefulManager = GracefulShutdownManager()
+        await TaskLocals.$gracefulShutdownManager.withValue(shutdownGracefulManager) {
+            await withThrowingTaskGroup(of: Void.self) { group in
+                group.addTask {
+                    var cont: AsyncStream<Void>.Continuation!
+                    let stream = AsyncStream<Void> { cont = $0 }
+                    let continuation = cont!
+
+                    await withShutdownGracefulHandler {
+                        await withTaskGroup(of: Void.self) { group in
+                            group.addTask {
+                                await stream.first { _ in true }
+                            }
+
+                            await group.waitForAll()
+                        }
+                    } onGracefulShutdown: {
+                        continuation.finish()
+                    }
+                }
+
+                await group.shutdownGracefullyAll()
+
+                try! await group.waitForAll()
+            }
+        }
+    }
 }