Added publish and close listeners

Author: adam-fowler

Sources/MQTTNIO/MQTTChannelHandlers.swift

@@ -45,7 +45,7 @@ struct ByteToMQTTMessageDecoder: ByteToMessageDecoder {
                 } catch MQTTSerializer.Error.incompletePacket {
                     return .needMoreData
                 } catch {
-                    self.client.publishCallback(.failure(error))
+                    self.client.publishListeners.notify(.failure(error))
                 }
                 return .continue
             case .CONNACK:
@@ -72,11 +72,11 @@ struct ByteToMQTTMessageDecoder: ByteToMessageDecoder {
     func publish(_ message: MQTTPublishMessage) {
         switch message.publish.qos {
         case .atMostOnce:
-            client.publishCallback(.success(message.publish))
+            client.publishListeners.notify(.success(message.publish))
         case .atLeastOnce:
             client.connection!.sendMessageNoWait(MQTTAckMessage(type: .PUBACK, packetId: message.packetId))
                 .map { _ in return message.publish }
-                .whenComplete { self.client.publishCallback($0) }
+                .whenComplete { self.client.publishListeners.notify($0) }
         case .exactlyOnce:
             client.connection!.sendMessage(MQTTAckMessage(type: .PUBREC, packetId: message.packetId)) { newMessage in
                 guard newMessage.packetId == message.packetId else { return false }
@@ -87,7 +87,7 @@ struct ByteToMQTTMessageDecoder: ByteToMessageDecoder {
                 self.client.connection!.sendMessageNoWait(MQTTAckMessage(type: .PUBCOMP, packetId: message.packetId))
             }
             .map { _ in return message.publish }
-            .whenComplete { self.client.publishCallback($0) }
+            .whenComplete { self.client.publishListeners.notify($0) }
         }
 
     }

Sources/MQTTNIO/MQTTClient.swift

@@ -32,8 +32,6 @@ final public class MQTTClient {
     var logger: Logger
     /// Client configuration
     let configuration: Configuration
-    /// Called whenever a publish event occurs
-    let publishCallback: (Result<MQTTPublishInfo, Swift.Error>) -> ()
 
     /// Connection client is using
     var connection: MQTTConnection?
@@ -90,8 +88,7 @@ final public class MQTTClient {
         port: Int? = nil,
         eventLoopGroupProvider: NIOEventLoopGroupProvider,
         logger: Logger? = nil,
-        configuration: Configuration = Configuration(),
-        publishCallback: @escaping (Result<MQTTPublishInfo, Swift.Error>) -> () = { _ in }
+        configuration: Configuration = Configuration()
     ) {
         self.host = host
         if let port = port {
@@ -109,7 +106,6 @@ final public class MQTTClient {
             }
         }
         self.configuration = configuration
-        self.publishCallback = publishCallback
         self.connection = nil
         self.logger = logger ?? Self.loggingDisabled
         self.eventLoopGroupProvider = eventLoopGroupProvider
@@ -153,7 +149,8 @@ final public class MQTTClient {
         return MQTTConnection.create(client: self, pingInterval: pingInterval)
             .flatMap { connection -> EventLoopFuture<MQTTInboundMessage> in
                 self.connection = connection
-                connection.closeFuture.whenComplete { _ in
+                connection.closeFuture.whenComplete { result in
+                    self.closeListeners.notify(result)
                     self.connection = nil
                 }
                 // attach client identifier to logger
@@ -259,6 +256,30 @@ final public class MQTTClient {
                 return future ?? self.eventLoopGroup.next().makeSucceededFuture(())
             }
     }
+    
+    /// Add named publish listener. Called whenever a PUBLISH message is received from the server
+    public func addPublishListener(named name: String, _ listener: @escaping (Result<MQTTPublishInfo, Swift.Error>) -> ()) {
+        publishListeners.addListener(named: name, listener: listener)
+    }
+    
+    /// Remove named publish listener
+    public func removePublishListener(named name: String) {
+        publishListeners.removeListener(named: name)
+    }
+    
+    /// Add close listener. Called whenever the connection is closed
+    public func addCloseListener(named name: String, _ listener: @escaping (Result<Void, Swift.Error>) -> ()) {
+        closeListeners.addListener(named: name, listener: listener)
+    }
+    
+    /// Remove named close listener
+    public func removeCloseListener(named name: String) {
+        closeListeners.removeListener(named: name)
+    }
+    
+
+    var publishListeners = MQTTListeners<MQTTPublishInfo>()
+    var closeListeners = MQTTListeners<Void>()
 }
 
 extension Logger {

Sources/MQTTNIO/MQTTListeners.swift

@@ -0,0 +1,30 @@
+import NIO
+import NIOConcurrencyHelpers
+
+struct MQTTListeners<ReturnType> {
+    typealias Listener = (Result<ReturnType, Error>) -> ()
+    
+    
+    func notify(_ result: Result<ReturnType, Error>) {
+        lock.withLock {
+            listeners.values.forEach { listener in
+                listener(result)
+            }
+        }
+    }
+    
+    mutating func addListener(named name: String, listener: @escaping Listener) {
+        lock.withLock {
+            listeners[name] = listener
+        }
+    }
+    
+    mutating func removeListener(named name: String) {
+        lock.withLock {
+            listeners[name] = nil
+        }
+    }
+    
+    private let lock = Lock()
+    private var listeners: [String: Listener] = [:]
+}

Tests/MQTTNIOTests/MQTTNIOTests.swift

@@ -41,26 +41,24 @@ final class MQTTNIOTests: XCTestCase {
     -----END CERTIFICATE-----
     """
 
-    func createClient(cb: @escaping (Result<MQTTPublishInfo, Swift.Error>) -> () = { _ in }) -> MQTTClient {
+    func createClient() -> MQTTClient {
         MQTTClient(
             host: "test.mosquitto.org",
             port: 1883,
             eventLoopGroupProvider: .createNew,
-            logger: self.logger,
-            publishCallback: cb
+            logger: self.logger
         )
     }
-    func createWebSocketClient(cb: @escaping (Result<MQTTPublishInfo, Swift.Error>) -> () = { _ in }) -> MQTTClient {
+    func createWebSocketClient() -> MQTTClient {
         MQTTClient(
             host: "test.mosquitto.org",
             port: 8080,
             eventLoopGroupProvider: .createNew,
             logger: self.logger,
-            configuration: .init(useWebSockets: true, webSocketURLPath: "/mqtt"),
-            publishCallback: cb
+            configuration: .init(useWebSockets: true, webSocketURLPath: "/mqtt")
         )
     }
-    func createSSLClient(cb: @escaping (Result<MQTTPublishInfo, Swift.Error>) -> () = { _ in }) throws -> MQTTClient {
+    func createSSLClient() throws -> MQTTClient {
         let rootCertificate = try NIOSSLCertificate.fromPEMBytes([UInt8](mosquittoCertificate.utf8))
         let tlsConfiguration: TLSConfiguration? = TLSConfiguration.forClient(
             trustRoots: .certificates(rootCertificate)
@@ -70,8 +68,7 @@ final class MQTTNIOTests: XCTestCase {
             port: 8883,
             eventLoopGroupProvider: .createNew,
             logger: self.logger,
-            configuration: .init(useSSL: true, tlsConfiguration: tlsConfiguration),
-            publishCallback: cb
+            configuration: .init(useSSL: true, tlsConfiguration: tlsConfiguration)
         )
     }
 
@@ -178,7 +175,8 @@ final class MQTTNIOTests: XCTestCase {
         )
         let client = self.createWebSocketClient()
         try connect(to: client, identifier: "soto_publisher")
-        let client2 = self.createWebSocketClient() { result in
+        let client2 = self.createWebSocketClient()
+        client2.addPublishListener(named: "test") { result in
             switch result {
             case .success(let publish):
                 var buffer = publish.payload