From a3422d7576a849a669772eea6153f62dfad3527e Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Wed, 21 Feb 2024 21:24:14 +0100 Subject: [PATCH 1/5] Add test case --- .../GracefulShutdownTests.swift | 84 +++++++++++++++++++ 1 file changed, 84 insertions(+) diff --git a/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift b/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift index 90085e6..6645490 100644 --- a/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift +++ b/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift @@ -353,4 +353,88 @@ final class GracefulShutdownTests: XCTestCase { group.cancelAll() } } + + func testCancelOnGracefulShutdownSurvivesCancellation() async throws { + await withTaskGroup(of: Void.self) { group in + group.addTask { + await withGracefulShutdownHandler { + await cancelOnGracefulShutdown { + await OnlyCancellationWaiter().cancellation + + try! await uncancellable { + try! await Task.sleep(for: .milliseconds(500)) + } + } + } onGracefulShutdown: { + XCTFail("Unexpect graceful shutdown") + } + } + + group.cancelAll() + } + } + + func testCancelOnGracefulShutdownSurvivesErrorThrown() async throws { + struct MyError: Error, Equatable {} + + await withTaskGroup(of: Void.self) { group in + group.addTask { + do { + try await withGracefulShutdownHandler { + try await cancelOnGracefulShutdown { + await OnlyCancellationWaiter().cancellation + + try! await uncancellable { + try! await Task.sleep(for: .milliseconds(500)) + } + + throw MyError() + } + } onGracefulShutdown: { + XCTFail("Unexpect graceful shutdown") + } + XCTFail("Expected to have thrown") + } catch { + XCTAssertEqual(error as? MyError, MyError()) + } + } + + group.cancelAll() + } + } +} + +func uncancellable(_ closure: @escaping @Sendable () async throws -> ()) async throws { + let task = Task { + try await closure() + } + + try await task.value +} + +private actor OnlyCancellationWaiter { + private var taskContinuation: CheckedContinuation? + + @usableFromInline + init() {} + + @usableFromInline + var cancellation: Void { + get async { + await withTaskCancellationHandler { + await withCheckedContinuation { continuation in + self.taskContinuation = continuation + } + } onCancel: { + Task { + await self.finish() + } + } + } + } + + private func finish() { + self.taskContinuation?.resume() + self.taskContinuation = nil + } } From ea2715c26912251762847ec25861d57accd1dd57 Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Wed, 21 Feb 2024 21:25:19 +0100 Subject: [PATCH 2/5] Fix tests --- .../AsyncGracefulShutdownSequence.swift | 7 ++--- .../ServiceLifecycle/CancellationWaiter.swift | 28 ++++++++++--------- .../ServiceLifecycle/GracefulShutdown.swift | 28 +++++++++++++------ 3 files changed, 38 insertions(+), 25 deletions(-) diff --git a/Sources/ServiceLifecycle/AsyncGracefulShutdownSequence.swift b/Sources/ServiceLifecycle/AsyncGracefulShutdownSequence.swift index 653aad1..209776c 100644 --- a/Sources/ServiceLifecycle/AsyncGracefulShutdownSequence.swift +++ b/Sources/ServiceLifecycle/AsyncGracefulShutdownSequence.swift @@ -20,7 +20,7 @@ @usableFromInline struct AsyncGracefulShutdownSequence: AsyncSequence, Sendable { @usableFromInline - typealias Element = Void + typealias Element = CancellationWaiter.Reason @inlinable init() {} @@ -36,9 +36,8 @@ struct AsyncGracefulShutdownSequence: AsyncSequence, Sendable { init() {} @inlinable - func next() async throws -> Element? { - try await CancellationWaiter().wait() - return () + func next() async -> Element? { + await CancellationWaiter().wait() } } } diff --git a/Sources/ServiceLifecycle/CancellationWaiter.swift b/Sources/ServiceLifecycle/CancellationWaiter.swift index d09a215..0bc73b2 100644 --- a/Sources/ServiceLifecycle/CancellationWaiter.swift +++ b/Sources/ServiceLifecycle/CancellationWaiter.swift @@ -15,36 +15,38 @@ /// An actor that provides a function to wait on cancellation/graceful shutdown. @usableFromInline actor CancellationWaiter { - private var taskContinuation: CheckedContinuation? + @usableFromInline + enum Reason { + case cancelled + case gracefulShutdown + } + + private var taskContinuation: CheckedContinuation? @usableFromInline init() {} @usableFromInline - func wait() async throws { - try await withTaskCancellationHandler { - try await withGracefulShutdownHandler { - try await withCheckedThrowingContinuation { continuation in + func wait() async -> Reason { + await withTaskCancellationHandler { + await withGracefulShutdownHandler { + await withCheckedContinuation { (continuation: CheckedContinuation) in self.taskContinuation = continuation } } onGracefulShutdown: { Task { - await self.finish() + await self.finish(reason: .gracefulShutdown) } } } onCancel: { Task { - await self.finish(throwing: CancellationError()) + await self.finish(reason: .cancelled) } } } - private func finish(throwing error: Error? = nil) { - if let error { - self.taskContinuation?.resume(throwing: error) - } else { - self.taskContinuation?.resume() - } + private func finish(reason: Reason) { + self.taskContinuation?.resume(returning: reason) self.taskContinuation = nil } } diff --git a/Sources/ServiceLifecycle/GracefulShutdown.swift b/Sources/ServiceLifecycle/GracefulShutdown.swift index 9c9c745..b3bcf37 100644 --- a/Sources/ServiceLifecycle/GracefulShutdown.swift +++ b/Sources/ServiceLifecycle/GracefulShutdown.swift @@ -95,13 +95,21 @@ public func withTaskCancellationOrGracefulShutdownHandler( /// /// - Throws: `CancellationError` if the task is cancelled. public func gracefulShutdown() async throws { - try await AsyncGracefulShutdownSequence().first { _ in true } + switch await AsyncGracefulShutdownSequence().first(where: { _ in true }) { + case .cancelled: + throw CancellationError() + case .gracefulShutdown: + return + case .none: + fatalError() + } } /// This is just a helper type for the result of our task group. enum ValueOrGracefulShutdown: Sendable { case value(T) case gracefulShutdown + case cancelled } /// Cancels the closure when a graceful shutdown was triggered. @@ -115,11 +123,15 @@ public func cancelOnGracefulShutdown(_ operation: @Sendable @escapi } group.addTask { - for try await _ in AsyncGracefulShutdownSequence() { - return .gracefulShutdown + for await reason in AsyncGracefulShutdownSequence() { + switch reason { + case .cancelled: + return .cancelled + case .gracefulShutdown: + return .gracefulShutdown + } } - - throw CancellationError() + fatalError("Unexpectedly didn't exit the task before") } let result = try await group.next() @@ -128,13 +140,13 @@ public func cancelOnGracefulShutdown(_ operation: @Sendable @escapi switch result { case .value(let t): return t - case .gracefulShutdown: + + case .gracefulShutdown, .cancelled: switch try await group.next() { case .value(let t): return t - case .gracefulShutdown: + case .gracefulShutdown, .cancelled: fatalError("Unexpectedly got gracefulShutdown from group.next()") - case nil: fatalError("Unexpectedly got nil from group.next()") } From 79ccd7e7801ce8a31dab97b69880d7e7c26809a3 Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Thu, 22 Feb 2024 11:59:13 +0100 Subject: [PATCH 3/5] Make code simpler --- Sources/ServiceLifecycle/GracefulShutdown.swift | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/Sources/ServiceLifecycle/GracefulShutdown.swift b/Sources/ServiceLifecycle/GracefulShutdown.swift index b3bcf37..e455087 100644 --- a/Sources/ServiceLifecycle/GracefulShutdown.swift +++ b/Sources/ServiceLifecycle/GracefulShutdown.swift @@ -123,15 +123,12 @@ public func cancelOnGracefulShutdown(_ operation: @Sendable @escapi } group.addTask { - for await reason in AsyncGracefulShutdownSequence() { - switch reason { - case .cancelled: - return .cancelled - case .gracefulShutdown: - return .gracefulShutdown - } + switch await CancellationWaiter().wait() { + case .cancelled: + return .cancelled + case .gracefulShutdown: + return .gracefulShutdown } - fatalError("Unexpectedly didn't exit the task before") } let result = try await group.next() From 1360779eef91bdfe5526159beaf0c3807baa2e0e Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Thu, 22 Feb 2024 12:23:24 +0100 Subject: [PATCH 4/5] Fix Sendable --- Sources/ServiceLifecycle/CancellationWaiter.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Sources/ServiceLifecycle/CancellationWaiter.swift b/Sources/ServiceLifecycle/CancellationWaiter.swift index 0bc73b2..1476fec 100644 --- a/Sources/ServiceLifecycle/CancellationWaiter.swift +++ b/Sources/ServiceLifecycle/CancellationWaiter.swift @@ -16,7 +16,7 @@ @usableFromInline actor CancellationWaiter { @usableFromInline - enum Reason { + enum Reason: Sendable { case cancelled case gracefulShutdown } From 140acf2ce3274ff282077f958e320a96987191ce Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Thu, 22 Feb 2024 12:25:02 +0100 Subject: [PATCH 5/5] swift-format --- Tests/ServiceLifecycleTests/GracefulShutdownTests.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift b/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift index 6645490..d71892b 100644 --- a/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift +++ b/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift @@ -404,7 +404,7 @@ final class GracefulShutdownTests: XCTestCase { } } -func uncancellable(_ closure: @escaping @Sendable () async throws -> ()) async throws { +func uncancellable(_ closure: @escaping @Sendable () async throws -> Void) async throws { let task = Task { try await closure() }