add semantics for stopping tasks that have no start handler (#77)

motivation: address shutdown for tasks which allocate resources on construction. before this change, only tasks that have been started are shutdown, this creates an issue when the resources are allocated on construction since the item may not get started at all (for example due to error in earlier task) and thus will never be shutdown and the resources will leak.

changes: add sematics to request shutdown even if not started, and assume this sematics when registering with start: .none or with registerShutdown

Author: tomerd

Sources/Lifecycle/Lifecycle.swift

@@ -26,29 +26,39 @@ import Logging
 /// Represents an item that can be started and shut down
 public protocol LifecycleTask {
     var label: String { get }
+    var shutdownIfNotStarted: Bool { get }
     func start(_ callback: @escaping (Error?) -> Void)
     func shutdown(_ callback: @escaping (Error?) -> Void)
 }
 
+extension LifecycleTask {
+    public var shutdownIfNotStarted: Bool {
+        return false
+    }
+}
+
 // MARK: - LifecycleHandler
 
 /// Supported startup and shutdown method styles
 public struct LifecycleHandler {
-    private let body: (@escaping (Error?) -> Void) -> Void
+    public typealias Callback = (@escaping (Error?) -> Void) -> Void
+
+    private let body: Callback?
 
     /// Initialize a `LifecycleHandler` based on a completion handler.
     ///
     /// - parameters:
     ///    - callback: the underlying completion handler
-    public init(_ callback: @escaping (@escaping (Error?) -> Void) -> Void) {
+    ///    - noop: the underlying completion handler is a no-op
+    public init(_ callback: Callback?) {
         self.body = callback
     }
 
     /// Asynchronous `LifecycleHandler` based on a completion handler.
     ///
     /// - parameters:
     ///    - callback: the underlying completion handler
-    public static func async(_ callback: @escaping (@escaping (Error?) -> Void) -> Void) -> LifecycleHandler {
+    public static func async(_ callback: @escaping Callback) -> LifecycleHandler {
         return LifecycleHandler(callback)
     }
 
@@ -69,13 +79,18 @@ public struct LifecycleHandler {
 
     /// Noop `LifecycleHandler`.
     public static var none: LifecycleHandler {
-        return LifecycleHandler { callback in
+        return LifecycleHandler(nil)
+    }
+
+    internal func run(_ callback: @escaping (Error?) -> Void) {
+        let body = self.body ?? { callback in
             callback(nil)
         }
+        body(callback)
     }
 
-    internal func run(_ callback: @escaping (Error?) -> Void) {
-        self.body(callback)
+    internal var noop: Bool {
+        return self.body == nil
     }
 }
 
@@ -242,12 +257,9 @@ public class ComponentLifecycle: LifecycleTask {
     private let logger: Logger
     internal let shutdownGroup = DispatchGroup()
 
-    private var state = State.idle
+    private var state = State.idle([])
     private let stateLock = Lock()
 
-    private var tasks = [LifecycleTask]()
-    private let tasksLock = Lock()
-
     /// Creates a `ComponentLifecycle` instance.
     ///
     /// - parameters:
@@ -275,7 +287,9 @@ public class ComponentLifecycle: LifecycleTask {
     ///    - on: `DispatchQueue` to run the handlers callback  on
     ///    - callback: The handler which is called after the start operation completes. The parameter will be `nil` on success and contain the `Error` otherwise.
     public func start(on queue: DispatchQueue, _ callback: @escaping (Error?) -> Void) {
-        let tasks = self.tasksLock.withLock { self.tasks }
+        guard case .idle(let tasks) = (self.stateLock.withLock { self.state }) else {
+            preconditionFailure("invalid state, \(self.state)")
+        }
         self._start(on: queue, tasks: tasks, callback: callback)
     }
 
@@ -328,7 +342,6 @@ public class ComponentLifecycle: LifecycleTask {
             self.stateLock.unlock()
             setupShutdownListener(queue)
         case .started(let queue, let tasks):
-            self.state = .shuttingDown(queue)
             self.stateLock.unlock()
             setupShutdownListener(queue)
             self._shutdown(on: queue, tasks: tasks, callback: self.shutdownGroup.leave)
@@ -362,7 +375,12 @@ public class ComponentLifecycle: LifecycleTask {
             case .shuttingDown:
                 self.stateLock.unlock()
                 // shutdown was called while starting, or start failed, shutdown what we can
-                let stoppable = started < tasks.count ? Array(tasks.prefix(started + 1)) : tasks
+                let stoppable: [LifecycleTask]
+                if started < tasks.count {
+                    stoppable = tasks.enumerated().filter { $0.offset <= started || $0.element.shutdownIfNotStarted }.map { $0.element }
+                } else {
+                    stoppable = tasks
+                }
                 self._shutdown(on: queue, tasks: stoppable) {
                     callback(error)
                     self.shutdownGroup.leave()
@@ -443,7 +461,7 @@ public class ComponentLifecycle: LifecycleTask {
     }
 
     private enum State {
-        case idle
+        case idle([LifecycleTask])
         case starting(DispatchQueue)
         case started(DispatchQueue, [LifecycleTask])
         case shuttingDown(DispatchQueue)
@@ -454,12 +472,10 @@ public class ComponentLifecycle: LifecycleTask {
 extension ComponentLifecycle: LifecycleTasksContainer {
     public func register(_ tasks: [LifecycleTask]) {
         self.stateLock.withLock {
-            guard case .idle = self.state else {
+            guard case .idle(let existing) = self.state else {
                 preconditionFailure("invalid state, \(self.state)")
             }
-        }
-        self.tasksLock.withLock {
-            self.tasks.append(contentsOf: tasks)
+            self.state = .idle(existing + tasks)
         }
     }
 }
@@ -489,7 +505,7 @@ extension LifecycleTasksContainer {
     ///    - start: `Handler` to perform the startup.
     ///    - shutdown: `Handler` to perform the shutdown.
     public func register(label: String, start: LifecycleHandler, shutdown: LifecycleHandler) {
-        self.register(_LifecycleTask(label: label, start: start, shutdown: shutdown))
+        self.register(_LifecycleTask(label: label, shutdownIfNotStarted: nil, start: start, shutdown: shutdown))
     }
 
     /// Adds a `LifecycleTask` to a `LifecycleTasks` collection.
@@ -504,9 +520,17 @@ extension LifecycleTasksContainer {
 
 internal struct _LifecycleTask: LifecycleTask {
     let label: String
+    let shutdownIfNotStarted: Bool
     let start: LifecycleHandler
     let shutdown: LifecycleHandler
 
+    init(label: String, shutdownIfNotStarted: Bool? = nil, start: LifecycleHandler, shutdown: LifecycleHandler) {
+        self.label = label
+        self.shutdownIfNotStarted = shutdownIfNotStarted ?? start.noop
+        self.start = start
+        self.shutdown = shutdown
+    }
+
     func start(_ callback: @escaping (Error?) -> Void) {
         self.start.run(callback)
     }

Sources/LifecycleNIOCompat/Bridge.swift

@@ -47,13 +47,13 @@ extension LifecycleHandler {
     }
 }
 
-extension ServiceLifecycle {
+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:
-    func start(on eventLoop: EventLoop) -> EventLoopFuture<Void> {
+    public 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+XCTest.swift

@@ -54,6 +54,9 @@ extension ComponentLifecycleTests {
             ("testNIOFailure", testNIOFailure),
             ("testInternalState", testInternalState),
             ("testExternalState", testExternalState),
+            ("testNOOPHandlers", testNOOPHandlers),
+            ("testShutdownOnlyStarted", testShutdownOnlyStarted),
+            ("testShutdownWhenStartFailedIfAsked", testShutdownWhenStartFailedIfAsked),
         ]
     }
 }

Tests/LifecycleTests/ComponentLifecycleTests.swift

@@ -82,6 +82,7 @@ final class ComponentLifecycleTests: XCTestCase {
         let items = (1 ... Int.random(in: 10 ... 20)).map { index -> LifecycleTask in
             let id = "item-\(index)"
             return _LifecycleTask(label: id,
+                                  shutdownIfNotStarted: false,
                                   start: .sync {
                                       dispatchPrecondition(condition: .onQueue(testQueue))
                                       startCalls.append(id)
@@ -839,4 +840,148 @@ final class ComponentLifecycleTests: XCTestCase {
         lifecycle.wait()
         XCTAssertEqual(state, .shutdown, "expected item to be shutdown, but \(state)")
     }
+
+    func testNOOPHandlers() {
+        let none = LifecycleHandler.none
+        XCTAssertEqual(none.noop, true)
+
+        let sync = LifecycleHandler.sync {}
+        XCTAssertEqual(sync.noop, false)
+
+        let async = LifecycleHandler.async { _ in }
+        XCTAssertEqual(async.noop, false)
+
+        let custom = LifecycleHandler { _ in }
+        XCTAssertEqual(custom.noop, false)
+    }
+
+    func testShutdownOnlyStarted() {
+        class Item {
+            let label: String
+            let sempahore: DispatchSemaphore
+            let failStart: Bool
+            let exptectedState: State
+            var state = State.idle
+
+            deinit {
+                XCTAssertEqual(self.state, self.exptectedState, "\"\(self.label)\" should be \(self.exptectedState)")
+                self.sempahore.signal()
+            }
+
+            init(label: String, failStart: Bool, exptectedState: State, sempahore: DispatchSemaphore) {
+                self.label = label
+                self.failStart = failStart
+                self.exptectedState = exptectedState
+                self.sempahore = sempahore
+            }
+
+            func start() throws {
+                self.state = .started
+                if self.failStart {
+                    self.state = .error
+                    throw InitError()
+                }
+            }
+
+            func shutdown() throws {
+                self.state = .shutdown
+            }
+
+            enum State {
+                case idle
+                case started
+                case shutdown
+                case error
+            }
+
+            struct InitError: Error {}
+        }
+
+        let count = Int.random(in: 10 ..< 20)
+        let sempahore = DispatchSemaphore(value: count)
+        let lifecycle = ServiceLifecycle(configuration: .init(shutdownSignal: nil))
+
+        for index in 0 ..< count {
+            let item = Item(label: "\(index)", failStart: index == count / 2, exptectedState: index <= count / 2 ? .shutdown : .idle, sempahore: sempahore)
+            lifecycle.register(label: item.label, start: .sync(item.start), shutdown: .sync(item.shutdown))
+        }
+
+        lifecycle.start { error in
+            XCTAssertNotNil(error, "expecting error")
+            lifecycle.shutdown()
+        }
+        lifecycle.wait()
+
+        XCTAssertEqual(.success, sempahore.wait(timeout: .now() + 1))
+    }
+
+    func testShutdownWhenStartFailedIfAsked() {
+        class DestructionSensitive {
+            let label: String
+            let failStart: Bool
+            let sempahore: DispatchSemaphore
+            var state = State.idle
+
+            deinit {
+                XCTAssertEqual(self.state, .shutdown, "\"\(self.label)\" should be shutdown")
+                self.sempahore.signal()
+            }
+
+            init(label: String, failStart: Bool = false, sempahore: DispatchSemaphore) {
+                self.label = label
+                self.failStart = failStart
+                self.sempahore = sempahore
+            }
+
+            func start() throws {
+                self.state = .started
+                if self.failStart {
+                    self.state = .error
+                    throw InitError()
+                }
+            }
+
+            func shutdown() throws {
+                self.state = .shutdown
+            }
+
+            enum State {
+                case idle
+                case started
+                case shutdown
+                case error
+            }
+
+            struct InitError: Error {}
+        }
+
+        let sempahore = DispatchSemaphore(value: 6)
+        let lifecycle = ServiceLifecycle(configuration: .init(shutdownSignal: nil))
+
+        let item1 = DestructionSensitive(label: "1", sempahore: sempahore)
+        lifecycle.register(label: item1.label, start: .sync(item1.start), shutdown: .sync(item1.shutdown))
+
+        let item2 = DestructionSensitive(label: "2", sempahore: sempahore)
+        lifecycle.registerShutdown(label: item2.label, .sync(item2.shutdown))
+
+        let item3 = DestructionSensitive(label: "3", failStart: true, sempahore: sempahore)
+        lifecycle.register(label: item3.label, start: .sync(item3.start), shutdown: .sync(item3.shutdown))
+
+        let item4 = DestructionSensitive(label: "4", sempahore: sempahore)
+        lifecycle.registerShutdown(label: item4.label, .sync(item4.shutdown))
+
+        let item5 = DestructionSensitive(label: "5", sempahore: sempahore)
+        lifecycle.register(label: item5.label, start: .none, shutdown: .sync(item5.shutdown))
+
+        let item6 = DestructionSensitive(label: "6", sempahore: sempahore)
+        lifecycle.register(_LifecycleTask(label: item6.label, shutdownIfNotStarted: true, start: .sync(item6.start), shutdown: .sync(item6.shutdown)))
+
+        lifecycle.start { error in
+            XCTAssertNotNil(error, "expecting error")
+            lifecycle.shutdown()
+        }
+        lifecycle.wait()
+
+        XCTAssertEqual(.success, sempahore.wait(timeout: .now() + 1))
+    }
 }