From d7fe0e731499a8dcce53bf4cbbc812c8e565d3a7 Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Thu, 22 Feb 2024 12:43:51 +0100 Subject: [PATCH] `cancelOnGracefulShutdown` hangs, if cancellation is not immediately (#177) * Add test case * Fix tests * Make code simpler * Fix Sendable * swift-format --- .../AsyncGracefulShutdownSequence.swift | 7 +- .../ServiceLifecycle/CancellationWaiter.swift | 28 ++++--- .../ServiceLifecycle/GracefulShutdown.swift | 23 +++-- .../GracefulShutdownTests.swift | 84 +++++++++++++++++++ 4 files changed, 118 insertions(+), 24 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..1476fec 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: Sendable { + 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..e455087 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,12 @@ public func cancelOnGracefulShutdown(_ operation: @Sendable @escapi } group.addTask { - for try await _ in AsyncGracefulShutdownSequence() { + switch await CancellationWaiter().wait() { + case .cancelled: + return .cancelled + case .gracefulShutdown: return .gracefulShutdown } - - throw CancellationError() } let result = try await group.next() @@ -128,13 +137,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()") } diff --git a/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift b/Tests/ServiceLifecycleTests/GracefulShutdownTests.swift index 90085e6..d71892b 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 -> Void) 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 + } }