shutdown() should cancel the signal handlers installed by start()

motivation: allow easier testing of shutdown hooks

changes:
* introduce ServiceLifecycle.removeTrap which removes a trap
* call ServiceLifecycle.removeTrap when setting up the shutdown hook
* make the shutdown hook cleanup into a lifecycle task to ensure correct ordering
* add tests
* improve logging

rdar://89552798

Author: tomerd

Sources/Lifecycle/Lifecycle.swift

@@ -30,12 +30,22 @@ public protocol LifecycleTask {
     var shutdownIfNotStarted: Bool { get }
     func start(_ callback: @escaping (Error?) -> Void)
     func shutdown(_ callback: @escaping (Error?) -> Void)
+    var logStart: Bool { get }
+    var logShutdown: Bool { get }
 }
 
-extension LifecycleTask {
-    public var shutdownIfNotStarted: Bool {
+public extension LifecycleTask {
+    var shutdownIfNotStarted: Bool {
         return false
     }
+
+    var logStart: Bool {
+        return true
+    }
+
+    var logShutdown: Bool {
+        return true
+    }
 }
 
 // MARK: - LifecycleHandler
@@ -97,8 +107,8 @@ public struct LifecycleHandler {
 
 #if canImport(_Concurrency) && compiler(>=5.5.2)
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
-extension LifecycleHandler {
-    public init(_ handler: @escaping () async throws -> Void) {
+public extension LifecycleHandler {
+    init(_ handler: @escaping () async throws -> Void) {
         self = LifecycleHandler { callback in
             Task {
                 do {
@@ -111,7 +121,7 @@ extension LifecycleHandler {
         }
     }
 
-    public static func async(_ handler: @escaping () async throws -> Void) -> LifecycleHandler {
+    static func async(_ handler: @escaping () async throws -> Void) -> LifecycleHandler {
         return LifecycleHandler(handler)
     }
 }
@@ -161,8 +171,8 @@ public struct LifecycleStartHandler<State> {
 
 #if canImport(_Concurrency) && compiler(>=5.5.2)
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
-extension LifecycleStartHandler {
-    public init(_ handler: @escaping () async throws -> State) {
+public extension LifecycleStartHandler {
+    init(_ handler: @escaping () async throws -> State) {
         self = LifecycleStartHandler { callback in
             Task {
                 do {
@@ -175,7 +185,7 @@ extension LifecycleStartHandler {
         }
     }
 
-    public static func async(_ handler: @escaping () async throws -> State) -> LifecycleStartHandler {
+    static func async(_ handler: @escaping () async throws -> State) -> LifecycleStartHandler {
         return LifecycleStartHandler(handler)
     }
 }
@@ -223,8 +233,8 @@ public struct LifecycleShutdownHandler<State> {
 
 #if canImport(_Concurrency) && compiler(>=5.5.2)
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
-extension LifecycleShutdownHandler {
-    public init(_ handler: @escaping (State) async throws -> Void) {
+public extension LifecycleShutdownHandler {
+    init(_ handler: @escaping (State) async throws -> Void) {
         self = LifecycleShutdownHandler { state, callback in
             Task {
                 do {
@@ -237,7 +247,7 @@ extension LifecycleShutdownHandler {
         }
     }
 
-    public static func async(_ handler: @escaping (State) async throws -> Void) -> LifecycleShutdownHandler {
+    static func async(_ handler: @escaping (State) async throws -> Void) -> LifecycleShutdownHandler {
         return LifecycleShutdownHandler(handler)
     }
 }
@@ -317,9 +327,14 @@ public struct ServiceLifecycle {
                 self.log("intercepted signal: \(signal)")
                 self.shutdown()
             }, cancelAfterTrap: true)
-            self.underlying.shutdownGroup.notify(queue: .global()) {
-                signalSource.cancel()
-            }
+            // register cleanup as the last task
+            self.registerShutdown(label: "\(signal) shutdown hook cleanup", .sync {
+                // cancel if not already canceled by the trap
+                if !signalSource.isCancelled {
+                    signalSource.cancel()
+                    ServiceLifecycle.removeTrap(signal: signal)
+                }
+            })
         }
     }
 
@@ -328,7 +343,7 @@ public struct ServiceLifecycle {
     }
 }
 
-extension ServiceLifecycle {
+public extension ServiceLifecycle {
     private static var trapped: Set<Int32> = []
     private static let trappedLock = Lock()
 
@@ -340,27 +355,39 @@ extension ServiceLifecycle {
     ///    - on: DispatchQueue to run the signal handler on (default global dispatch queue)
     ///    - cancelAfterTrap: Defaults to false, which means the signal handler can be run multiple times. If true, the DispatchSignalSource will be cancelled after being trapped once.
     /// - returns: a `DispatchSourceSignal` for the given trap. The source must be cancelled by the caller.
-    public static func trap(signal sig: Signal, handler: @escaping (Signal) -> Void, on queue: DispatchQueue = .global(), cancelAfterTrap: Bool = false) -> DispatchSourceSignal {
+    static func trap(signal sig: Signal, handler: @escaping (Signal) -> Void, on queue: DispatchQueue = .global(), cancelAfterTrap: Bool = false) -> DispatchSourceSignal {
         // on linux, we can call singal() once per process
         self.trappedLock.withLockVoid {
-            if !trapped.contains(sig.rawValue) {
+            if !self.trapped.contains(sig.rawValue) {
                 signal(sig.rawValue, SIG_IGN)
-                trapped.insert(sig.rawValue)
+                self.trapped.insert(sig.rawValue)
             }
         }
         let signalSource = DispatchSource.makeSignalSource(signal: sig.rawValue, queue: queue)
         signalSource.setEventHandler {
+            // run handler first
+            handler(sig)
+            // then cancel trap if so requested
             if cancelAfterTrap {
                 signalSource.cancel()
+                self.removeTrap(signal: sig)
             }
-            handler(sig)
         }
         signalSource.resume()
         return signalSource
     }
 
+    static func removeTrap(signal sig: Signal) {
+        self.trappedLock.withLockVoid {
+            if self.trapped.contains(sig.rawValue) {
+                signal(sig.rawValue, SIG_DFL)
+                self.trapped.remove(sig.rawValue)
+            }
+        }
+    }
+
     /// A system signal
-    public struct Signal: Equatable, CustomStringConvertible {
+    struct Signal: Equatable, CustomStringConvertible {
         internal var rawValue: CInt
 
         public static let TERM = Signal(rawValue: SIGTERM)
@@ -395,9 +422,9 @@ extension ServiceLifecycle: LifecycleTasksContainer {
     }
 }
 
-extension ServiceLifecycle {
+public extension ServiceLifecycle {
     /// `ServiceLifecycle` configuration options.
-    public struct Configuration {
+    struct Configuration {
         /// Defines the `label` for the lifeycle and its Logger
         public var label: String
         /// Defines the `Logger` to log with.
@@ -413,7 +440,8 @@ extension ServiceLifecycle {
                     logger: Logger? = nil,
                     callbackQueue: DispatchQueue = .global(),
                     shutdownSignal: [Signal]? = [.TERM, .INT],
-                    installBacktrace: Bool = true) {
+                    installBacktrace: Bool = true)
+        {
             self.label = label
             self.logger = logger ?? Logger(label: label)
             self.callbackQueue = callbackQueue
@@ -433,7 +461,7 @@ struct ShutdownError: Error {
 public class ComponentLifecycle: LifecycleTask {
     public let label: String
     fileprivate let logger: Logger
-    internal let shutdownGroup = DispatchGroup()
+    fileprivate let shutdownGroup = DispatchGroup()
 
     private var state = State.idle([])
     private let stateLock = Lock()
@@ -596,13 +624,15 @@ public class ComponentLifecycle: LifecycleTask {
 
     private func startTask(on queue: DispatchQueue, tasks: [LifecycleTask], index: Int, callback: @escaping ([LifecycleTask], Error?) -> Void) {
         // async barrier
-        let start = { (callback) -> Void in queue.async { tasks[index].start(callback) } }
-        let callback = { (index, error) -> Void in queue.async { callback(index, error) } }
+        let start = { callback in queue.async { tasks[index].start(callback) } }
+        let callback = { index, error in queue.async { callback(index, error) } }
 
         if index >= tasks.count {
             return callback(tasks, nil)
         }
-        self.logger.info("starting tasks [\(tasks[index].label)]")
+        if tasks[index].logStart {
+            self.logger.info("starting tasks [\(tasks[index].label)]")
+        }
         let startTime = DispatchTime.now()
         start { error in
             Timer(label: "\(self.label).\(tasks[index].label).lifecycle.start").recordNanoseconds(DispatchTime.now().uptimeNanoseconds - startTime.uptimeNanoseconds)
@@ -642,14 +672,16 @@ public class ComponentLifecycle: LifecycleTask {
 
     private func shutdownTask(on queue: DispatchQueue, tasks: [LifecycleTask], index: Int, errors: [String: Error]?, callback: @escaping ([String: Error]?) -> Void) {
         // async barrier
-        let shutdown = { (callback) -> Void in queue.async { tasks[index].shutdown(callback) } }
-        let callback = { (errors) -> Void in queue.async { callback(errors) } }
+        let shutdown = { callback in queue.async { tasks[index].shutdown(callback) } }
+        let callback = { errors in queue.async { callback(errors) } }
 
         if index >= tasks.count {
             return callback(errors)
         }
 
-        self.logger.info("stopping tasks [\(tasks[index].label)]")
+        if tasks[index].logShutdown {
+            self.logger.info("stopping tasks [\(tasks[index].label)]")
+        }
         let startTime = DispatchTime.now()
         shutdown { error in
             Timer(label: "\(self.label).\(tasks[index].label).lifecycle.shutdown").recordNanoseconds(DispatchTime.now().uptimeNanoseconds - startTime.uptimeNanoseconds)
@@ -694,12 +726,12 @@ public protocol LifecycleTasksContainer {
     func register(_ tasks: [LifecycleTask])
 }
 
-extension LifecycleTasksContainer {
+public extension LifecycleTasksContainer {
     /// Adds a `LifecycleTask` to a `LifecycleTasks` collection.
     ///
     /// - parameters:
     ///    - tasks: one or more `LifecycleTask`.
-    public func register(_ tasks: LifecycleTask ...) {
+    func register(_ tasks: LifecycleTask ...) {
         self.register(tasks)
     }
 
@@ -709,7 +741,7 @@ extension LifecycleTasksContainer {
     ///    - label: label of the item, useful for debugging.
     ///    - start: `Handler` to perform the startup.
     ///    - shutdown: `Handler` to perform the shutdown.
-    public func register(label: String, start: LifecycleHandler, shutdown: LifecycleHandler, shutdownIfNotStarted: Bool? = nil) {
+    func register(label: String, start: LifecycleHandler, shutdown: LifecycleHandler, shutdownIfNotStarted: Bool? = nil) {
         self.register(_LifecycleTask(label: label, shutdownIfNotStarted: shutdownIfNotStarted, start: start, shutdown: shutdown))
     }
 
@@ -718,7 +750,7 @@ extension LifecycleTasksContainer {
     /// - parameters:
     ///    - label: label of the item, useful for debugging.
     ///    - handler: `Handler` to perform the shutdown.
-    public func registerShutdown(label: String, _ handler: LifecycleHandler) {
+    func registerShutdown(label: String, _ handler: LifecycleHandler) {
         self.register(label: label, start: .none, shutdown: handler)
     }
 
@@ -728,7 +760,7 @@ extension LifecycleTasksContainer {
     ///    - label: label of the item, useful for debugging.
     ///    - start: `LifecycleStartHandler` to perform the startup and return the state.
     ///    - shutdown: `LifecycleShutdownHandler` to perform the shutdown given the state.
-    public func registerStateful<State>(label: String, start: LifecycleStartHandler<State>, shutdown: LifecycleShutdownHandler<State>) {
+    func registerStateful<State>(label: String, start: LifecycleStartHandler<State>, shutdown: LifecycleShutdownHandler<State>) {
         self.register(StatefulLifecycleTask(label: label, start: start, shutdown: shutdown))
     }
 }
@@ -739,12 +771,16 @@ internal struct _LifecycleTask: LifecycleTask {
     let shutdownIfNotStarted: Bool
     let start: LifecycleHandler
     let shutdown: LifecycleHandler
+    let logStart: Bool
+    let logShutdown: Bool
 
     init(label: String, shutdownIfNotStarted: Bool? = nil, start: LifecycleHandler, shutdown: LifecycleHandler) {
         self.label = label
         self.shutdownIfNotStarted = shutdownIfNotStarted ?? start.noop
         self.start = start
         self.shutdown = shutdown
+        self.logStart = !start.noop
+        self.logShutdown = !shutdown.noop
     }
 
     func start(_ callback: @escaping (Error?) -> Void) {

Sources/Lifecycle/Locks.swift

@@ -81,7 +81,7 @@ extension Lock {
     /// - Parameter body: The block to execute while holding the lock.
     /// - Returns: The value returned by the block.
     @inlinable
-    internal func withLock<T>(_ body: () throws -> T) rethrows -> T {
+    func withLock<T>(_ body: () throws -> T) rethrows -> T {
         self.lock()
         defer {
             self.unlock()
@@ -91,7 +91,7 @@ extension Lock {
 
     // specialise Void return (for performance)
     @inlinable
-    internal func withLockVoid(_ body: () throws -> Void) rethrows {
+    func withLockVoid(_ body: () throws -> Void) rethrows {
         try self.withLock(body)
     }
 }

Sources/LifecycleNIOCompat/Bridge.swift

@@ -15,12 +15,12 @@
 import Lifecycle
 import NIO
 
-extension LifecycleHandler {
+public extension LifecycleHandler {
     /// Asynchronous `LifecycleHandler` based on an `EventLoopFuture`.
     ///
     /// - parameters:
     ///    - future: function returning the underlying `EventLoopFuture`
-    public static func eventLoopFuture(_ future: @escaping () -> EventLoopFuture<Void>) -> LifecycleHandler {
+    static func eventLoopFuture(_ future: @escaping () -> EventLoopFuture<Void>) -> LifecycleHandler {
         return LifecycleHandler { callback in
             future().whenComplete { result in
                 switch result {
@@ -34,13 +34,13 @@ extension LifecycleHandler {
     }
 }
 
-extension LifecycleHandler {
+public extension LifecycleHandler {
     /// `LifecycleHandler` that cancels a `RepeatedTask`.
     ///
     /// - parameters:
     ///    - task: `RepeatedTask` to be cancelled
     ///    - on: `EventLoop` to use for cancelling the task
-    public static func cancelRepeatedTask(_ task: RepeatedTask, on eventLoop: EventLoop) -> LifecycleHandler {
+    static func cancelRepeatedTask(_ task: RepeatedTask, on eventLoop: EventLoop) -> LifecycleHandler {
         return self.eventLoopFuture {
             let promise = eventLoop.makePromise(of: Void.self)
             task.cancel(promise: promise)
@@ -49,12 +49,12 @@ extension LifecycleHandler {
     }
 }
 
-extension LifecycleStartHandler {
+public extension LifecycleStartHandler {
     /// Asynchronous `LifecycleStartHandler` based on an `EventLoopFuture`.
     ///
     /// - parameters:
     ///    - future: function returning the underlying `EventLoopFuture`
-    public static func eventLoopFuture(_ future: @escaping () -> EventLoopFuture<State>) -> LifecycleStartHandler {
+    static func eventLoopFuture(_ future: @escaping () -> EventLoopFuture<State>) -> LifecycleStartHandler {
         return LifecycleStartHandler { callback in
             future().whenComplete { result in
                 callback(result)
@@ -63,12 +63,12 @@ extension LifecycleStartHandler {
     }
 }
 
-extension LifecycleShutdownHandler {
+public extension LifecycleShutdownHandler {
     /// Asynchronous `LifecycleShutdownHandler` based on an `EventLoopFuture`.
     ///
     /// - parameters:
     ///    - future: function returning the underlying `EventLoopFuture`
-    public static func eventLoopFuture(_ future: @escaping (State) -> EventLoopFuture<Void>) -> LifecycleShutdownHandler {
+    static func eventLoopFuture(_ future: @escaping (State) -> EventLoopFuture<Void>) -> LifecycleShutdownHandler {
         return LifecycleShutdownHandler { state, callback in
             future(state).whenComplete { result in
                 switch result {
@@ -82,13 +82,13 @@ extension LifecycleShutdownHandler {
     }
 }
 
-extension ComponentLifecycle {
+public extension ComponentLifecycle {
     /// Starts the provided `LifecycleItem` array.
     /// Startup is performed in the order of items provided.
     ///
     /// - parameters:
     ///    - eventLoop: The `eventLoop` which is used to generate the `EventLoopFuture` that is returned. After the start the future is fulfilled:
-    public func start(on eventLoop: EventLoop) -> EventLoopFuture<Void> {
+    func start(on eventLoop: EventLoop) -> EventLoopFuture<Void> {
         let promise = eventLoop.makePromise(of: Void.self)
         self.start { error in
             if let error = error {

Tests/LifecycleTests/ComponentLifecycleTests.swift

@@ -53,7 +53,7 @@ final class ComponentLifecycleTests: XCTestCase {
                                       dispatchPrecondition(condition: .onQueue(.global()))
                                       XCTAssertTrue(startCalls.contains(id))
                                       stopCalls.append(id)
-                                 })
+                                  })
         }
         lifecycle.register(items)
 
@@ -92,7 +92,7 @@ final class ComponentLifecycleTests: XCTestCase {
                                       dispatchPrecondition(condition: .onQueue(testQueue))
                                       XCTAssertTrue(startCalls.contains(id))
                                       stopCalls.append(id)
-                                 })
+                                  })
         }
         lifecycle.register(items)