add support for stateful handler

motivation: allow lifecycle handlers to have the library manage the state for them so they do not need to do that manually

changes:
* introduce LifecycleStartHandler and LifecycleShutdownHandler which can handle state on behalf of the lifecycle item
* add registerStateful function to regsiter stateful handlers
* add tests

TODO: update docs

Author: tomerd

Sources/Lifecycle/Lifecycle.swift

@@ -41,25 +41,27 @@ extension LifecycleTask {
 
 /// Supported startup and shutdown method styles
 public struct LifecycleHandler {
+    @available(*, deprecated, renamed: "Handler")
     public typealias Callback = (@escaping (Error?) -> Void) -> Void
 
-    private let body: Callback?
+    public typealias Handler = (@escaping (Error?) -> Void) -> Void
+
+    private let underlying: Handler?
 
     /// Initialize a `LifecycleHandler` based on a completion handler.
     ///
     /// - parameters:
-    ///    - callback: the underlying completion handler
-    ///    - noop: the underlying completion handler is a no-op
-    public init(_ callback: Callback?) {
-        self.body = callback
+    ///    - handler: the underlying completion handler
+    public init(_ handler: Handler?) {
+        self.underlying = handler
     }
 
     /// Asynchronous `LifecycleHandler` based on a completion handler.
     ///
     /// - parameters:
-    ///    - callback: the underlying completion handler
-    public static func async(_ callback: @escaping Callback) -> LifecycleHandler {
-        return LifecycleHandler(callback)
+    ///    - handler: the underlying async handler
+    public static func async(_ handler: @escaping Handler) -> LifecycleHandler {
+        return LifecycleHandler(handler)
     }
 
     /// Asynchronous `LifecycleHandler` based on a blocking, throwing function.
@@ -82,15 +84,101 @@ public struct LifecycleHandler {
         return LifecycleHandler(nil)
     }
 
-    internal func run(_ callback: @escaping (Error?) -> Void) {
-        let body = self.body ?? { callback in
+    internal func run(_ completionHandler: @escaping (Error?) -> Void) {
+        let body = self.underlying ?? { callback in
             callback(nil)
         }
-        body(callback)
+        body(completionHandler)
     }
 
     internal var noop: Bool {
-        return self.body == nil
+        return self.underlying == nil
+    }
+}
+
+// MARK: - Stateful Lifecycle Handlers
+
+/// LifecycleHandler for starting stateful tasks. The state can then be fed into a LifecycleShutdownHandler
+public struct LifecycleStartHandler<State> {
+    public typealias Handler = (@escaping (Result<State, Error>) -> Void) -> Void
+
+    private let underlying: Handler
+
+    /// Initialize a `LifecycleHandler` based on a completion handler.
+    ///
+    /// - parameters:
+    ///    - callback: the underlying completion handler
+    public init(_ handler: @escaping Handler) {
+        self.underlying = handler
+    }
+
+    /// Asynchronous `LifecycleStartHandler` based on a completion handler.
+    ///
+    /// - parameters:
+    ///    - handler: the underlying async handler
+    public static func async(_ handler: @escaping Handler) -> LifecycleStartHandler {
+        return LifecycleStartHandler(handler)
+    }
+
+    /// Asynchronous `LifecycleStartHandler` based on a blocking, throwing function.
+    ///
+    /// - parameters:
+    ///    - body: the underlying function
+    public static func sync(_ body: @escaping () throws -> State) -> LifecycleStartHandler {
+        return LifecycleStartHandler { completionHandler in
+            do {
+                let state = try body()
+                completionHandler(.success(state))
+            } catch {
+                completionHandler(.failure(error))
+            }
+        }
+    }
+
+    internal func run(_ completionHandler: @escaping (Result<State, Error>) -> Void) {
+        self.underlying(completionHandler)
+    }
+}
+
+/// LifecycleHandler for shutting down stateful tasks. The state comes from a LifecycleStartHandler
+public struct LifecycleShutdownHandler<State> {
+    public typealias Handler = (State, @escaping (Error?) -> Void) -> Void
+
+    private let underlying: Handler
+
+    /// Initialize a `LifecycleShutdownHandler` based on a completion handler.
+    ///
+    /// - parameters:
+    ///    - handler: the underlying completion handler
+    public init(_ handler: @escaping Handler) {
+        self.underlying = handler
+    }
+
+    /// Asynchronous `LifecycleShutdownHandler` based on a completion handler.
+    ///
+    /// - parameters:
+    ///    - handler: the underlying async handler
+    public static func async(_ handler: @escaping Handler) -> LifecycleShutdownHandler {
+        return LifecycleShutdownHandler(handler)
+    }
+
+    /// Asynchronous `LifecycleShutdownHandler` based on a blocking, throwing function.
+    ///
+    /// - parameters:
+    ///    - body: the underlying function
+    public static func sync(_ body: @escaping (State) throws -> Void) -> LifecycleShutdownHandler {
+        return LifecycleShutdownHandler { state, completionHandler in
+            do {
+                try body(state)
+                completionHandler(nil)
+            } catch {
+                completionHandler(error)
+            }
+        }
+    }
+
+    internal func run(state: State, _ completionHandler: @escaping (Error?) -> Void) {
+        self.underlying(state, completionHandler)
     }
 }
 
@@ -516,6 +604,16 @@ extension LifecycleTasksContainer {
     public func registerShutdown(label: String, _ handler: LifecycleHandler) {
         self.register(label: label, start: .none, shutdown: handler)
     }
+
+    /// Adds a  stateful `LifecycleTask` to a `LifecycleTasks` collection.
+    ///
+    /// - parameters:
+    ///    - 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>) {
+        self.register(StatefulLifecycleTask(label: label, start: start, shutdown: shutdown))
+    }
 }
 
 internal struct _LifecycleTask: LifecycleTask {
@@ -539,3 +637,42 @@ internal struct _LifecycleTask: LifecycleTask {
         self.shutdown.run(callback)
     }
 }
+
+internal class StatefulLifecycleTask<State>: LifecycleTask {
+    let label: String
+    let shutdownIfNotStarted: Bool = false
+    let start: LifecycleStartHandler<State>
+    let shutdown: LifecycleShutdownHandler<State>
+
+    let stateLock = Lock()
+    var state: State?
+
+    init(label: String, start: LifecycleStartHandler<State>, shutdown: LifecycleShutdownHandler<State>) {
+        self.label = label
+        self.start = start
+        self.shutdown = shutdown
+    }
+
+    func start(_ callback: @escaping (Error?) -> Void) {
+        self.start.run { result in
+            switch result {
+            case .failure(let error):
+                callback(error)
+            case .success(let state):
+                self.stateLock.withLock {
+                    self.state = state
+                }
+                callback(nil)
+            }
+        }
+    }
+
+    func shutdown(_ callback: @escaping (Error?) -> Void) {
+        guard let state = (self.stateLock.withLock { self.state }) else {
+            return callback(UnknownState())
+        }
+        self.shutdown.run(state: state, callback)
+    }
+
+    struct UnknownState: Error {}
+}

Sources/LifecycleNIOCompat/Bridge.swift

@@ -16,7 +16,7 @@ import Lifecycle
 import NIO
 
 extension LifecycleHandler {
-    /// Asynchronous `Lifecycle.Handler` based on an `EventLoopFuture`.
+    /// Asynchronous `LifecycleHandler` based on an `EventLoopFuture`.
     ///
     /// - parameters:
     ///    - future: function returning the underlying `EventLoopFuture`
@@ -32,8 +32,10 @@ extension LifecycleHandler {
             }
         }
     }
+}
 
-    /// `Lifecycle.Handler` that cancels a `RepeatedTask`.
+extension LifecycleHandler {
+    /// `LifecycleHandler` that cancels a `RepeatedTask`.
     ///
     /// - parameters:
     ///    - task: `RepeatedTask` to be cancelled
@@ -47,6 +49,39 @@ extension LifecycleHandler {
     }
 }
 
+extension LifecycleStartHandler {
+    /// Asynchronous `LifecycleStartHandler` based on an `EventLoopFuture`.
+    ///
+    /// - parameters:
+    ///    - future: function returning the underlying `EventLoopFuture`
+    public static func eventLoopFuture(_ future: @escaping () -> EventLoopFuture<State>) -> LifecycleStartHandler {
+        return LifecycleStartHandler { callback in
+            future().whenComplete { result in
+                callback(result)
+            }
+        }
+    }
+}
+
+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 {
+        return LifecycleShutdownHandler { state, callback in
+            future(state).whenComplete { result in
+                switch result {
+                case .success:
+                    callback(nil)
+                case .failure(let error):
+                    callback(error)
+                }
+            }
+        }
+    }
+}
+
 extension ComponentLifecycle {
     /// Starts the provided `LifecycleItem` array.
     /// Startup is performed in the order of items provided.

Tests/LifecycleTests/ComponentLifecycleTests+XCTest.swift

@@ -57,6 +57,15 @@ extension ComponentLifecycleTests {
             ("testNOOPHandlers", testNOOPHandlers),
             ("testShutdownOnlyStarted", testShutdownOnlyStarted),
             ("testShutdownWhenStartFailedIfAsked", testShutdownWhenStartFailedIfAsked),
+            ("testStatefulSync", testStatefulSync),
+            ("testStatefulSyncStartError", testStatefulSyncStartError),
+            ("testStatefulSyncShutdownError", testStatefulSyncShutdownError),
+            ("testStatefulAsync", testStatefulAsync),
+            ("testStatefulAsyncStartError", testStatefulAsyncStartError),
+            ("testStatefulAsyncShutdownError", testStatefulAsyncShutdownError),
+            ("testStatefulNIO", testStatefulNIO),
+            ("testStatefulNIOStartFailure", testStatefulNIOStartFailure),
+            ("testStatefulNIOShutdownFailure", testStatefulNIOShutdownFailure),
         ]
     }
 }