Async shutdown (#73)

Author: adam-fowler

Sources/MQTTNIO/AsyncAwaitSupport/MQTTClient+async.swift

@@ -5,6 +5,18 @@ import NIOCore
 
 @available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *)
 extension MQTTClient {
+    public func shutdown(queue: DispatchQueue = .global()) async throws {
+        return try await withUnsafeThrowingContinuation { cont in
+            shutdown(queue: queue) { error in
+                if let error = error {
+                    cont.resume(throwing: error)
+                } else {
+                    cont.resume()
+                }
+            }
+        }
+    }
+
     /// Connect to MQTT server
     ///
     /// Completes when CONNACK is received
@@ -111,7 +123,7 @@ public class MQTTPublishListener: AsyncSequence {
     }
 
     public __consuming func makeAsyncIterator() -> AsyncStream<Element>.AsyncIterator {
-        return stream.makeAsyncIterator()
+        return self.stream.makeAsyncIterator()
     }
 }
 

Sources/MQTTNIO/AsyncAwaitSupport/MQTTClientV5+async.swift

@@ -130,9 +130,8 @@ public class MQTTPublishIdListener: AsyncSequence {
     }
 
     public __consuming func makeAsyncIterator() -> AsyncStream<Element>.AsyncIterator {
-        return stream.makeAsyncIterator()
+        return self.stream.makeAsyncIterator()
     }
 }
 
-
 #endif // compiler(>=5.5)

Sources/MQTTNIO/MQTTClient.swift

@@ -65,6 +65,8 @@ public final class MQTTClient {
     private static let loggingDisabled = Logger(label: "MQTT-do-not-log", factory: { _ in SwiftLogNoOpLogHandler() })
     /// inflight messages
     private var inflight: MQTTInflight
+    /// flag to tell is client is shutdown
+    private let isShutdown = NIOAtomic<Bool>.makeAtomic(value: false)
 
     /// Create MQTT client
     /// - Parameters:
@@ -121,15 +123,88 @@ public final class MQTTClient {
         self.inflight = .init()
     }
 
-    /// Close down client. Must be called before the client is destroyed
+    deinit {
+        guard isShutdown.load() else {
+            preconditionFailure("Client not shut down before the deinit. Please call client.syncShutdown() when no longer needed.")
+        }
+    }
+
+    /// Shutdown client synchronously. Before an `MQTTClient` is deleted you need to call this function or the async version `shutdown`
+    /// to do a clean shutdown of the client. It closes the connection, notifies everything listening for shutdown and shuts down the
+    /// EventLoopGroup if the client created it
+    ///
+    /// - Throws: MQTTError.alreadyShutdown: You have already shutdown the client
     public func syncShutdownGracefully() throws {
-        try self.connection?.close().wait()
+        if let eventLoop = MultiThreadedEventLoopGroup.currentEventLoop {
+            preconditionFailure("""
+            BUG DETECTED: syncShutdown() must not be called when on an EventLoop.
+            Calling syncShutdown() on any EventLoop can lead to deadlocks.
+            Current eventLoop: \(eventLoop)
+            """)
+        }
+        let errorStorageLock = Lock()
+        var errorStorage: Error?
+        let continuation = DispatchWorkItem {}
+        self.shutdown(queue: DispatchQueue(label: "mqtt-client.shutdown")) { error in
+            if let error = error {
+                errorStorageLock.withLock {
+                    errorStorage = error
+                }
+            }
+            continuation.perform()
+        }
+        continuation.wait()
+        try errorStorageLock.withLock {
+            if let error = errorStorage {
+                throw error
+            }
+        }
+    }
+
+    /// Shutdown MQTTClient asynchronously. Before an `AWSClient` is deleted you need to call this function or the synchronous
+    /// version `syncShutdownGracefully` to do a clean shutdown of the client. It closes the connection, notifies everything
+    /// listening for shutdown and shuts down the EventLoopGroup if the client created it
+    ///
+    /// - Parameters:
+    ///   - queue: Dispatch Queue to run shutdown on
+    ///   - callback: Callback called when shutdown is complete. If there was an error it will return with Error in callback
+    public func shutdown(queue: DispatchQueue = .global(), _ callback: @escaping (Error?) -> Void) {
+        guard self.isShutdown.compareAndExchange(expected: false, desired: true) else {
+            callback(MQTTError.alreadyShutdown)
+            return
+        }
+        let eventLoop = self.eventLoopGroup.next()
+        let closeFuture: EventLoopFuture<Void>
         self.shutdownListeners.notify(.success(()))
+        if let connection = self.connection {
+            closeFuture = connection.close()
+        } else {
+            closeFuture = eventLoop.makeSucceededVoidFuture()
+        }
+        closeFuture.whenComplete { result in
+            let closeError: Error?
+            switch result {
+            case .failure(let error):
+                closeError = error
+            case .success:
+                closeError = nil
+            }
+            self.shutdownListeners.notify(.success(()))
+            self.shutdownEventLoopGroup(queue: queue) { error in
+                callback(closeError ?? error)
+            }
+        }
+    }
+
+    /// shutdown EventLoopGroup
+    private func shutdownEventLoopGroup(queue: DispatchQueue, _ callback: @escaping (Error?) -> Void) {
         switch self.eventLoopGroupProvider {
-        case .createNew:
-            try self.eventLoopGroup.syncShutdownGracefully()
         case .shared:
-            break
+            queue.async {
+                callback(nil)
+            }
+        case .createNew:
+            self.eventLoopGroup.shutdownGracefully(queue: queue, callback)
         }
     }
 

Sources/MQTTNIO/MQTTClientV5.swift

@@ -158,7 +158,8 @@ extension MQTTClient {
                 if case .success(let info) = result {
                     for property in info.properties {
                         if case .subscriptionIdentifier(let id) = property,
-                           id == subscriptionId {
+                           id == subscriptionId
+                        {
                             listener(info)
                             break
                         }

Sources/MQTTNIO/MQTTConnection.swift

@@ -181,7 +181,7 @@ final class MQTTConnection {
         let task = MQTTTask(on: channel.eventLoop, timeout: self.timeout, checkInbound: checkInbound)
         let taskHandler = MQTTTaskHandler(task: task, channel: channel)
 
-        channel.pipeline.addHandler(taskHandler, position: .before(self.unhandledHandler))
+        self.channel.pipeline.addHandler(taskHandler, position: .before(self.unhandledHandler))
             .flatMap {
                 self.channel.writeAndFlush(message)
             }

Sources/MQTTNIO/MQTTError.swift

@@ -20,6 +20,8 @@ public enum MQTTError: Error {
 
     /// You called connect on a client that is already connected to the broker
     case alreadyConnected
+    /// Client has already been shutdown
+    case alreadyShutdown
     /// We received an unexpected message while connecting
     case failedToConnect
     /// We received an unsuccessful connection return value

Sources/MQTTNIO/WebSocketHandler.swift

@@ -25,7 +25,7 @@ final class WebSocketHandler: ChannelDuplexHandler {
         guard context.channel.isActive else { return }
 
         let buffer = unwrapOutboundIn(data)
-        send(context: context, buffer: buffer, opcode: .binary, fin: true, promise: promise)
+        self.send(context: context, buffer: buffer, opcode: .binary, fin: true, promise: promise)
     }
 
     /// Read WebSocket frame

Tests/MQTTNIOTests/MQTTNIOTests+async.swift

@@ -51,11 +51,14 @@ final class AsyncMQTTNIOTests: XCTestCase {
         self.XCTRunAsyncAndBlock {
             try await client.connect()
             try await client.disconnect()
-            try client.syncShutdownGracefully()
+            try await client.shutdown()
         }
     }
 
     func testPublishSubscribe() {
+        let expectation = XCTestExpectation(description: "testPublishSubscribe")
+        expectation.expectedFulfillmentCount = 1
+
         let client = self.createClient(identifier: "testPublish+async")
         let client2 = self.createClient(identifier: "testPublish+async2")
         let payloadString = "Hello"
@@ -69,15 +72,19 @@ final class AsyncMQTTNIOTests: XCTestCase {
                     var buffer = publish.payload
                     let string = buffer.readString(length: buffer.readableBytes)
                     XCTAssertEqual(string, payloadString)
+                    expectation.fulfill()
                 case .failure(let error):
                     XCTFail("\(error)")
                 }
             }
             try await client.publish(to: "TestSubject", payload: ByteBufferAllocator().buffer(string: payloadString), qos: .atLeastOnce)
+
+            self.wait(for: [expectation], timeout: 2)
+
             try await client.disconnect()
-            Thread.sleep(forTimeInterval: 2)
             try await client2.disconnect()
-            try client.syncShutdownGracefully()
+            try await client.shutdown()
+            try await client2.shutdown()
         }
     }
 
@@ -95,18 +102,19 @@ final class AsyncMQTTNIOTests: XCTestCase {
             try await client.connect()
             try await client.ping()
             try await client.disconnect()
-            try client.syncShutdownGracefully()
+            try await client.shutdown()
         }
     }
 
     func testAsyncSequencePublishListener() {
-        let client = createClient(identifier: "testAsyncSequencePublishListener+async", version: .v5_0)
-        let client2 = createClient(identifier: "testAsyncSequencePublishListener+async2", version: .v5_0)
-        let payloadString = "Hello"
-        let expectation = XCTestExpectation(description: "publish listener")
+        let expectation = XCTestExpectation(description: "testAsyncSequencePublishListener")
         expectation.expectedFulfillmentCount = 3
-        
-        XCTRunAsyncAndBlock {
+
+        let client = self.createClient(identifier: "testAsyncSequencePublishListener+async", version: .v5_0)
+        let client2 = self.createClient(identifier: "testAsyncSequencePublishListener+async2", version: .v5_0)
+        let payloadString = "Hello"
+
+        self.XCTRunAsyncAndBlock {
             try await client.connect()
             try await client2.connect()
             _ = try await client2.v5.subscribe(to: [.init(topicFilter: "TestSubject", qos: .atLeastOnce)])
@@ -131,24 +139,24 @@ final class AsyncMQTTNIOTests: XCTestCase {
             Thread.sleep(forTimeInterval: 0.5)
             try await client2.disconnect()
             Thread.sleep(forTimeInterval: 0.5)
-            try client.syncShutdownGracefully()
-            try client2.syncShutdownGracefully()
+            try await client.shutdown()
+            try await client2.shutdown()
 
             _ = await task.result
         }
         wait(for: [expectation], timeout: 5.0)
     }
 
     func testAsyncSequencePublishSubscriptionIdListener() {
-        let client = createClient(identifier: "testAsyncSequencePublishSubscriptionIdListener+async", version: .v5_0)
-        let client2 = createClient(identifier: "testAsyncSequencePublishSubscriptionIdListener+async2", version: .v5_0)
+        let client = self.createClient(identifier: "testAsyncSequencePublishSubscriptionIdListener+async", version: .v5_0)
+        let client2 = self.createClient(identifier: "testAsyncSequencePublishSubscriptionIdListener+async2", version: .v5_0)
         let payloadString = "Hello"
         let expectation = XCTestExpectation(description: "publish listener")
         let expectation2 = XCTestExpectation(description: "publish listener2")
         expectation.expectedFulfillmentCount = 3
         expectation2.expectedFulfillmentCount = 2
 
-        XCTRunAsyncAndBlock {
+        self.XCTRunAsyncAndBlock {
             try await client.connect()
             try await client2.connect()
             _ = try await client2.v5.subscribe(to: [.init(topicFilter: "TestSubject", qos: .atLeastOnce)], properties: [.subscriptionIdentifier(1)])

Tests/MQTTNIOTests/MQTTNIOv5Tests.swift

@@ -234,6 +234,7 @@ final class MQTTNIOv5Tests: XCTestCase {
         connack = try client.v5.connect(cleanStart: false).wait()
         XCTAssertEqual(connack.sessionPresent, true)
         try client.disconnect().wait()
+        try client.syncShutdownGracefully()
     }
 
     func testPersistentSession() throws {