|
| 1 | +/// Execute an operation with a graceful shutdown handler that’s immediately invoked if the current task is shutting down gracefully. |
| 2 | +/// |
| 3 | +/// This doesn’t check for graceful shutdown, and always executes the passed operation. |
| 4 | +/// The operation executes on the calling execution context and does not suspend by itself, unless the code contained within the closure does. |
| 5 | +/// If graceful shutdown occurs while the operation is running, the graceful shutdown handler will execute concurrently with the operation. |
| 6 | +/// |
| 7 | +/// When `withShutdownGracefulHandler` is used in a Task that has already been gracefully shutdown, the `onGracefulShutdown` handler |
| 8 | +/// will be executed immediately before operation gets to execute. This allows the `onGracefulShutdown` handler to set some external “shutdown” flag |
| 9 | +/// that the operation may be atomically checking for in order to avoid performing any actual work once the operation gets to run. |
| 10 | +/// |
| 11 | +/// A common use-case is to listen to graceful shutdown and use the `ServerQuiescingHelper` from `swift-nio-extras` to |
| 12 | +/// trigger the quiescing sequence. Furthermore, graceful shutdown will propagate to any child task that is currently executing |
| 13 | +/// |
| 14 | +/// - Parameters: |
| 15 | +/// - operation: The actual operation. |
| 16 | +/// - handler: The handler which is invoked once graceful shutdown has been triggered. |
| 17 | +@_unsafeInheritExecutor |
| 18 | +public func withShutdownGracefulHandler<T>( |
| 19 | + operation: () async throws -> T, |
| 20 | + onGracefulShutdown handler: @Sendable @escaping () -> Void |
| 21 | +) async rethrows -> T { |
| 22 | + guard let gracefulShutdownManager = TaskLocals.gracefulShutdownManager else { |
| 23 | + print("Trying to setup a graceful shutdown handler 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 setup a shutdown graceful handler outside the ServiceRunner.run method. Not setting up the handler.") |
| 24 | + return try await operation() |
| 25 | + } |
| 26 | + |
| 27 | + // We have to keep track of our handler here to remove it once the operation is finished. |
| 28 | + let handlerNumber = await gracefulShutdownManager.registerHandler(handler) |
| 29 | + |
| 30 | + let result = try await operation() |
| 31 | + |
| 32 | + // Great the operation is finished. If we have a number we need to remove the handler. |
| 33 | + if let handlerNumber = handlerNumber { |
| 34 | + await gracefulShutdownManager.removeHandler(handlerNumber) |
| 35 | + } |
| 36 | + |
| 37 | + return result |
| 38 | +} |
| 39 | + |
| 40 | +@_spi(Testing) |
| 41 | +public enum TaskLocals { |
| 42 | + @TaskLocal |
| 43 | + @_spi(Testing) |
| 44 | + public static var gracefulShutdownManager: GracefulShutdownManager? |
| 45 | +} |
| 46 | + |
| 47 | +@_spi(Testing) |
| 48 | +public actor GracefulShutdownManager { |
| 49 | + /// The currently registered handlers. |
| 50 | + private var handlers = [(UInt64, () -> Void)]() |
| 51 | + /// A counter to assign a unique number to each handler. |
| 52 | + private var handlerCounter: UInt64 = 0 |
| 53 | + /// A boolean indicating if we have been shutdown already. |
| 54 | + private var isShuttingDown = false |
| 55 | + |
| 56 | + @_spi(Testing) |
| 57 | + public init() {} |
| 58 | + |
| 59 | + func registerHandler(_ handler: @Sendable @escaping () -> Void) -> UInt64? { |
| 60 | + if self.isShuttingDown { |
| 61 | + handler() |
| 62 | + return nil |
| 63 | + } else { |
| 64 | + defer { |
| 65 | + self.handlerCounter += 1 |
| 66 | + } |
| 67 | + let handlerNumber = self.handlerCounter |
| 68 | + self.handlers.append((handlerNumber, handler)) |
| 69 | + |
| 70 | + return handlerNumber |
| 71 | + } |
| 72 | + } |
| 73 | + |
| 74 | + func removeHandler(_ handlerNumber: UInt64) { |
| 75 | + self.handlers.removeAll { $0.0 == handlerNumber } |
| 76 | + } |
| 77 | + |
| 78 | + @_spi(Testing) |
| 79 | + public func shutdownGracefully() { |
| 80 | + guard !self.isShuttingDown else { |
| 81 | + fatalError("Tried to shutdown gracefully more than once") |
| 82 | + } |
| 83 | + self.isShuttingDown = true |
| 84 | + |
| 85 | + for handler in self.handlers { |
| 86 | + handler.1() |
| 87 | + } |
| 88 | + |
| 89 | + self.handlers.removeAll() |
| 90 | + } |
| 91 | +} |
0 commit comments