Skip to content

Commit 754bf5c

Browse files
committed
Remove shutdownGracefully method and switch to a handler based approach.
1 parent 5d4f7f3 commit 754bf5c

10 files changed

+421
-275
lines changed

.swiftformat

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,5 +19,7 @@
1919
--disable redundantReturn
2020
--disable preferKeyPath
2121
--disable sortedSwitchCases
22+
--disable hoistTry
23+
--disable hoistAwait
2224

2325
# rules

Sources/ServiceLifecycle/CancellableContinuation.swift

Lines changed: 0 additions & 37 deletions
This file was deleted.
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
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+
}

Sources/ServiceLifecycle/Service.swift

Lines changed: 0 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -27,25 +27,4 @@ public protocol Service: Sendable {
2727
/// - Handling incoming connections and requests
2828
/// - Background refreshes
2929
func run() async throws
30-
31-
/// This method is called when the ``ServiceRunner`` receives a graceful shutdown signal.
32-
///
33-
/// Concrete implementations can implement this if they support graceful shutdown.
34-
func shutdownGracefully() async throws
35-
}
36-
37-
extension Service {
38-
public var isLongRunning: Bool {
39-
// By default all services are treated as long running services.
40-
true
41-
}
42-
43-
public func run() async throws {
44-
// We are just going to suspend here until we get cancelled.
45-
try await CancellableContinuation().run()
46-
}
47-
48-
public func shutdownGracefully() async throws {
49-
// no-op
50-
}
5130
}

0 commit comments

Comments
 (0)