End publish listener on connection close if clean session is set to true

Author: adam-fowler

Sources/MQTTNIO/AsyncAwaitSupport/MQTTClient+async.swift

@@ -173,18 +173,25 @@ public class MQTTPublishListener: AsyncSequence {
         let name = UUID().uuidString
         self.client = client
         self.name = name
+        let cleanSession = client.connection?.cleanSession ?? true
         self.stream = AsyncStream { cont in
             client.addPublishListener(named: name) { result in
                 cont.yield(result)
             }
             client.addShutdownListener(named: name) { _ in
                 cont.finish()
             }
+            client.addCloseListener(named: name) { connectResult in
+                if cleanSession {
+                    cont.finish()
+                }
+            }
         }
     }
 
     deinit {
         self.client.removePublishListener(named: self.name)
+        self.client.removeCloseListener(named: self.name)
         self.client.removeShutdownListener(named: self.name)
     }
 

Sources/MQTTNIO/AsyncAwaitSupport/MQTTClientV5+async.swift

@@ -166,18 +166,25 @@ public class MQTTPublishIdListener: AsyncSequence {
         let name = UUID().uuidString
         self.client = client.client
         self.name = name
+        let cleanSession = client.client.connection?.cleanSession ?? true
         self.stream = AsyncStream { cont in
             client.addPublishListener(named: name, subscriptionId: subscriptionId) { result in
                 cont.yield(result)
             }
             client.client.addShutdownListener(named: name) { _ in
                 cont.finish()
             }
+            client.client.addCloseListener(named: name) { connectResult in
+                if cleanSession {
+                    cont.finish()
+                }
+            }
         }
     }
 
     deinit {
         self.client.removePublishListener(named: self.name)
+        self.client.removeCloseListener(named: self.name)
         self.client.removeShutdownListener(named: self.name)
     }
 

Sources/MQTTNIO/MQTTClient.swift

@@ -463,9 +463,9 @@ public final class MQTTClient {
     }
 
     var connectionParameters = ConnectionParameters()
-    let publishListeners = MQTTListeners<MQTTPublishInfo>()
-    let closeListeners = MQTTListeners<Void>()
-    let shutdownListeners = MQTTListeners<Void>()
+    let publishListeners = MQTTListeners<Result<MQTTPublishInfo, Error>>()
+    let closeListeners = MQTTListeners<Result<Void, Error>>()
+    let shutdownListeners = MQTTListeners<Result<Void, Error>>()
     private var _connection: MQTTConnection?
     private var lock = NIOLock()
 }
@@ -478,7 +478,7 @@ extension MQTTClient {
     ) -> EventLoopFuture<MQTTConnAckPacket> {
         let pingInterval = self.configuration.pingInterval ?? TimeAmount.seconds(max(Int64(packet.keepAliveSeconds - 5), 5))
 
-        let connectFuture = MQTTConnection.create(client: self, pingInterval: pingInterval)
+        let connectFuture = MQTTConnection.create(client: self, cleanSession: packet.cleanSession, pingInterval: pingInterval)
         let eventLoop = connectFuture.eventLoop
         return
             connectFuture

Sources/MQTTNIO/MQTTConnection.swift

@@ -30,19 +30,21 @@ import NIOSSL
 
 final class MQTTConnection {
     let channel: Channel
+    let cleanSession: Bool
     let timeout: TimeAmount?
     let taskHandler: MQTTTaskHandler
 
-    private init(channel: Channel, timeout: TimeAmount?, taskHandler: MQTTTaskHandler) {
+    private init(channel: Channel, cleanSession: Bool, timeout: TimeAmount?, taskHandler: MQTTTaskHandler) {
         self.channel = channel
+        self.cleanSession = cleanSession
         self.timeout = timeout
         self.taskHandler = taskHandler
     }
 
-    static func create(client: MQTTClient, pingInterval: TimeAmount) -> EventLoopFuture<MQTTConnection> {
+    static func create(client: MQTTClient, cleanSession: Bool, pingInterval: TimeAmount) -> EventLoopFuture<MQTTConnection> {
         let taskHandler = MQTTTaskHandler(client: client)
         return self.createBootstrap(client: client, pingInterval: pingInterval, taskHandler: taskHandler)
-            .map { MQTTConnection(channel: $0, timeout: client.configuration.timeout, taskHandler: taskHandler) }
+            .map { MQTTConnection(channel: $0, cleanSession: cleanSession, timeout: client.configuration.timeout, taskHandler: taskHandler) }
     }
 
     static func createBootstrap(client: MQTTClient, pingInterval: TimeAmount, taskHandler: MQTTTaskHandler) -> EventLoopFuture<Channel> {

Sources/MQTTNIO/MQTTListeners.swift

@@ -15,9 +15,9 @@ import NIO
 import NIOConcurrencyHelpers
 
 final class MQTTListeners<ReturnType> {
-    typealias Listener = (Result<ReturnType, Error>) -> Void
+    typealias Listener = (ReturnType) -> Void
 
-    func notify(_ result: Result<ReturnType, Error>) {
+    func notify(_ result: ReturnType) {
         let listeners = self.lock.withLock {
             return self.listeners
         }