From 9a538f941ca59f4fb45054080f173bacc3b92dc4 Mon Sep 17 00:00:00 2001
From: Adam Fowler <adamfowler71@gmail.com>
Date: Fri, 16 Feb 2024 12:00:26 +0000
Subject: [PATCH 1/2] Add withTaskCancellationOrGracefulShutdownHandler

---
 .../ServiceLifecycle/GracefulShutdown.swift   | 33 +++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/Sources/ServiceLifecycle/GracefulShutdown.swift b/Sources/ServiceLifecycle/GracefulShutdown.swift
index 76b8838..9c9c745 100644
--- a/Sources/ServiceLifecycle/GracefulShutdown.swift
+++ b/Sources/ServiceLifecycle/GracefulShutdown.swift
@@ -55,6 +55,39 @@ public func withGracefulShutdownHandler<T>(
     return try await operation()
 }
 
+/// Execute an operation with a graceful shutdown or task cancellation handler that’s immediately invoked if the current task is
+/// shutting down gracefully or has been cancelled.
+///
+/// This doesn’t check for graceful shutdown, and always executes the passed operation.
+/// The operation executes on the calling execution context and does not suspend by itself, unless the code contained within the closure does.
+/// If graceful shutdown or task cancellation occurs while the operation is running, the cancellation/graceful shutdown handler will execute
+/// concurrently with the operation.
+///
+/// When `withTaskCancellationOrGracefulShutdownHandler` is used in a Task that has already been gracefully shutdown or cancelled, the
+/// `onCancelOrGracefulShutdown` handler will be executed immediately before operation gets to execute. This allows the `onCancelOrGracefulShutdown`
+/// handler to set some external “shutdown” flag that the operation may be atomically checking for in order to avoid performing any actual work
+/// once the operation gets to run.
+///
+/// - Important: This method will only set up a graceful shutdown handler if run inside ``ServiceGroup`` otherwise no graceful shutdown handler
+/// will be set up.
+///
+/// - Parameters:
+///   - operation: The actual operation.
+///   - handler: The handler which is invoked once graceful shutdown or task cancellation has been triggered.
+// Unsafely inheriting the executor is safe to do here since we are not calling any other async method
+// except the operation. This makes sure no other executor hops would occur here.
+@_unsafeInheritExecutor
+public func withTaskCancellationOrGracefulShutdownHandler<T>(
+    operation: () async throws -> T,
+    onCancelOrGracefulShutdown handler: @Sendable @escaping () -> Void
+) async rethrows -> T {
+    return try await withTaskCancellationHandler {
+        try await withGracefulShutdownHandler(operation: operation, onGracefulShutdown: handler)
+    } onCancel: {
+        handler()
+    }
+}
+
 /// Waits until graceful shutdown is triggered.
 ///
 /// This method suspends the caller until graceful shutdown is triggered. If the calling task is cancelled before

From 1a0820f89f92405ce224cf9e4f8fcb0c966b2acd Mon Sep 17 00:00:00 2001
From: Adam Fowler <adamfowler71@gmail.com>
Date: Fri, 16 Feb 2024 12:00:47 +0000
Subject: [PATCH 2/2] Test withTaskCancellationOrGracefulShutdownHandler

---
 .../GracefulShutdownTests.swift               | 50 +++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift b/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift
index 13e082d..90085e6 100644
--- a/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift
+++ b/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift
@@ -303,4 +303,54 @@ final class GracefulShutdownTests: XCTestCase {
             XCTAssertTrue(error is CancellationError)
         }
     }
+
+    func testWithTaskCancellationOrGracefulShutdownHandler_GracefulShutdown() async throws {
+        var cont: AsyncStream<Void>.Continuation!
+        let stream = AsyncStream<Void> { cont = $0 }
+        let continuation = cont!
+
+        await testGracefulShutdown { gracefulShutdownTestTrigger in
+            await withTaskCancellationOrGracefulShutdownHandler {
+                await withTaskGroup(of: Void.self) { group in
+                    group.addTask {
+                        await stream.first { _ in true }
+                    }
+
+                    gracefulShutdownTestTrigger.triggerGracefulShutdown()
+
+                    await group.waitForAll()
+                }
+            } onCancelOrGracefulShutdown: {
+                continuation.finish()
+            }
+        }
+    }
+
+    func testWithTaskCancellationOrGracefulShutdownHandler_TaskCancellation() async throws {
+        var cont: AsyncStream<Void>.Continuation!
+        let stream = AsyncStream<Void> { cont = $0 }
+        let continuation = cont!
+
+        await withTaskGroup(of: Void.self) { group in
+            group.addTask {
+                var cancelCont: AsyncStream<CheckedContinuation<Void, Never>>.Continuation!
+                let cancelStream = AsyncStream<CheckedContinuation<Void, Never>> { cancelCont = $0 }
+                let cancelContinuation = cancelCont!
+                await withTaskCancellationOrGracefulShutdownHandler {
+                    await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
+                        cancelContinuation.yield(cont)
+                        continuation.finish()
+                    }
+                } onCancelOrGracefulShutdown: {
+                    Task {
+                        let cont = await cancelStream.first(where: { _ in true })!
+                        cont.resume()
+                    }
+                }
+            }
+            // wait for task to startup
+            _ = await stream.first { _ in true }
+            group.cancelAll()
+        }
+    }
 }