99// See CONTRIBUTORS.txt for the list of SwiftServiceLifecycle project authors
1010//
1111// SPDX-License-Identifier: Apache-2.0
12- //
12+
13+ import ConcurrencyHelpers
1314
1415/// Execute an operation with a graceful shutdown handler that’s immediately invoked if the current task is shutting down gracefully.
1516///
2526/// trigger the quiescing sequence. Furthermore, graceful shutdown will propagate to any child task that is currently executing
2627///
2728/// - Parameters:
29+ /// - requiresRunningInsideServiceRunner: Indicates if this method requires to be run from within a ``ServiceRunner`` child task.
30+ /// This defaults to `true` and if run outside of a ``ServiceRunner`` child task will `fatalError`. If set to `false` then
31+ /// no graceful shutdown handler will be setup if not called inside a ``ServiceRunner`` child task. This is useful for code that
32+ /// can run both inside and outside of``ServiceRunner`` child tasks.
2833/// - operation: The actual operation.
2934/// - handler: The handler which is invoked once graceful shutdown has been triggered.
35+ // Unsafely inheriting the executor is safe to do here since we are not calling any other async method
36+ // except the operation. This makes sure no other executor hops would occur here.
37+ @_unsafeInheritExecutor
3038public func withGracefulShutdownHandler< T> (
31- @_inheritActorContext operation: @Sendable ( ) async throws -> T ,
39+ requiresRunningInsideServiceRunner: Bool = true ,
40+ operation: ( ) async throws -> T ,
3241 onGracefulShutdown handler: @Sendable @escaping ( ) -> Void
3342) async rethrows -> T {
3443 guard let gracefulShutdownManager = TaskLocals . gracefulShutdownManager else {
35- print ( " WARNING: 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. " )
36- return try await operation ( )
44+ if !requiresRunningInsideServiceRunner {
45+ return try await operation ( )
46+ } else {
47+ fatalError ( " 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. " )
48+ }
3749 }
3850
3951 // We have to keep track of our handler here to remove it once the operation is finished.
40- let handlerID = await gracefulShutdownManager. registerHandler ( handler)
52+ let handlerID = gracefulShutdownManager. registerHandler ( handler)
53+ defer {
54+ if let handlerID = handlerID {
55+ gracefulShutdownManager. removeHandler ( handlerID)
56+ }
57+ }
4158
42- let result = try await operation ( )
59+ return try await operation ( )
60+ }
4361
44- // Great the operation is finished. If we have a number we need to remove the handler.
45- if let handlerID {
46- await gracefulShutdownManager. removeHandler ( handlerID)
47- }
62+ enum ValueOrGracefulShutdown < T> {
63+ case value( T )
64+ case gracefulShutdown
65+ }
66+
67+ /// Cancels the closure when a graceful shutdown was triggered.
68+ ///
69+ /// - Parameter operation: The actual operation.
70+ public func cancelOnGracefulShutdown< T> ( _ operation: @Sendable @escaping ( ) async throws -> T ) async rethrows -> T ? {
71+ return try await withThrowingTaskGroup ( of: ValueOrGracefulShutdown< T> . self ) { group in
72+ group. addTask {
73+ let value = try await operation ( )
74+ return . value( value)
75+ }
76+
77+ group. addTask {
78+ for await _ in AsyncGracefulShutdownSequence ( ) {
79+ return . gracefulShutdown
80+ }
81+
82+ throw CancellationError ( )
83+ }
4884
49- return result
85+ let result = try await group. next ( )
86+
87+ switch result {
88+ case . value( let t) :
89+ return t
90+ case . gracefulShutdown:
91+ group. cancelAll ( )
92+ switch try await group. next ( ) {
93+ case . value ( let t) :
94+ return t
95+ case . gracefulShutdown:
96+ fatalError ( " Unexpectedly got gracefulShutdown from group.next() " )
97+
98+ case nil:
99+ fatalError ( " Unexpectedly got nil from group.next() " )
100+ }
101+
102+ case nil:
103+ fatalError( " Unexpectedly got nil from group.next() " )
104+ }
105+ }
50106}
51107
52108@_spi( TestKit)
@@ -57,59 +113,70 @@ public enum TaskLocals {
57113}
58114
59115@_spi ( TestKit)
60- public actor GracefulShutdownManager {
116+ public final class GracefulShutdownManager : @ unchecked Sendable {
61117 struct Handler {
62118 /// The id of the handler.
63119 var id : UInt64
64120 /// The actual handler.
65121 var handler : ( ) -> Void
66122 }
67123
68- /// The currently registered handlers.
69- private var handlers = [ Handler] ( )
70- /// A counter to assign a unique number to each handler.
71- private var handlerCounter : UInt64 = 0
72- /// A boolean indicating if we have been shutdown already.
73- private var isShuttingDown = false
124+ struct State {
125+ /// The currently registered handlers.
126+ fileprivate var handlers = [ Handler] ( )
127+ /// A counter to assign a unique number to each handler.
128+ fileprivate var handlerCounter : UInt64 = 0
129+ /// A boolean indicating if we have been shutdown already.
130+ fileprivate var isShuttingDown = false
131+ }
132+
133+ private let state = LockedValueBox ( State ( ) )
74134
75135 @_spi ( TestKit)
76136 public init ( ) { }
77137
78138 func registerHandler( _ handler: @Sendable @escaping ( ) -> Void ) -> UInt64 ? {
79- if self . isShuttingDown {
80- handler ( )
81- return nil
82- } else {
83- defer {
84- self . handlerCounter += 1
139+ return self . state. withLockedValue { state in
140+ if state. isShuttingDown {
141+ // We are already shutting down so we just run the handler now.
142+ handler ( )
143+ return nil
144+ } else {
145+ defer {
146+ state. handlerCounter += 1
147+ }
148+ let handlerID = state. handlerCounter
149+ state. handlers. append ( . init( id: handlerID, handler: handler) )
150+
151+ return handlerID
85152 }
86- let handlerID = self . handlerCounter
87- self . handlers. append ( . init( id: handlerID, handler: handler) )
88-
89- return handlerID
90153 }
91154 }
92155
93156 func removeHandler( _ handlerID: UInt64 ) {
94- guard let index = self . handlers. firstIndex ( where: { $0. id == handlerID } ) else {
95- // This can happen because if shutdownGracefully ran while the operation was still in progress
96- return
97- }
157+ self . state. withLockedValue { state in
158+ guard let index = state. handlers. firstIndex ( where: { $0. id == handlerID } ) else {
159+ // This can happen because if shutdownGracefully ran while the operation was still in progress
160+ return
161+ }
98162
99- self . handlers. remove ( at: index)
163+ state. handlers. remove ( at: index)
164+ }
100165 }
101166
102167 @_spi ( TestKit)
103168 public func shutdownGracefully( ) {
104- guard !self . isShuttingDown else {
105- return
106- }
107- self . isShuttingDown = true
169+ self . state. withLockedValue { state in
170+ guard !state. isShuttingDown else {
171+ return
172+ }
173+ state. isShuttingDown = true
108174
109- for handler in self . handlers {
110- handler. handler ( )
111- }
175+ for handler in state . handlers {
176+ handler. handler ( )
177+ }
112178
113- self . handlers. removeAll ( )
179+ state. handlers. removeAll ( )
180+ }
114181 }
115182}
0 commit comments