feat: Renaming and moves some items around, listeners now manage reconnection events.

Sources/MQTTManager/Interface.swift

@@ -0,0 +1,197 @@
+import AsyncAlgorithms
+import Dependencies
+import DependenciesMacros
+import Foundation
+import Logging
+import MQTTNIO
+import NIO
+
+public extension DependencyValues {
+
+  /// A dependency that is responsible for managing the connection to
+  /// an MQTT broker.
+  var mqtt: MQTTManager {
+    get { self[MQTTManager.self] }
+    set { self[MQTTManager.self] = newValue }
+  }
+}
+
+/// Represents the interface needed for the ``MQTTConnectionService``.
+///
+/// See ``MQTTConnectionManagerLive`` module for live implementation.
+@DependencyClient
+public struct MQTTManager: Sendable {
+
+  public typealias ListenStream = AsyncStream<MQTTPublishInfo>
+
+  /// Connect to the MQTT broker.
+  public var connect: @Sendable () async throws -> Void
+
+  /// Create a stream of connection events.
+  public var connectionStream: @Sendable () throws -> AsyncStream<Event>
+
+  private var _listen: @Sendable ([String], MQTTQoS) async throws -> ListenStream
+
+  /// Publish a value to the MQTT broker for a given topic.
+  public var publish: @Sendable (PublishRequest) async throws -> Void
+
+  /// Shutdown the connection to the MQTT broker.
+  public var shutdown: @Sendable () -> Void
+
+  /// Perform an operation with the underlying MQTTClient, this can be useful in
+  /// tests, so this module needs imported with `@_spi(Testing) import` to use this method.
+  private var _withClient: @Sendable ((MQTTClient) async throws -> Void) async throws -> Void
+
+  /// Create an async stream that listens for changes to the given topics.
+  ///
+  /// - Parameters:
+  ///   - topics: The topics to listen for changes to.
+  ///   - qos: The MQTTQoS for the subscription.
+  public func listen(
+    to topics: [String],
+    qos: MQTTQoS = .atLeastOnce
+  ) async throws -> ListenStream {
+    try await _listen(topics, qos)
+  }
+
+  /// Create an async stream that listens for changes to the given topics.
+  ///
+  /// - Parameters:
+  ///   - topics: The topics to listen for changes to.
+  ///   - qos: The MQTTQoS for the subscription.
+  public func listen(
+    _ topics: String...,
+    qos: MQTTQoS = .atLeastOnce
+  ) async throws -> ListenStream {
+    try await listen(to: topics, qos: qos)
+  }
+
+  /// Publish a new value to the given topic.
+  ///
+  /// - Parameters:
+  ///   - payload: The value to publish.
+  ///   - topicName: The topic to publish the new value to.
+  ///   - qos: The MQTTQoS.
+  ///   - retain: The retain flag.
+  public func publish(
+    _ payload: ByteBuffer,
+    to topicName: String,
+    qos: MQTTQoS,
+    retain: Bool = false
+  ) async throws {
+    try await publish(.init(
+      topicName: topicName,
+      payload: payload,
+      qos: qos,
+      retain: retain
+    ))
+  }
+
+  /// Perform an operation with the underlying MQTTClient, this can be useful in
+  /// tests, so this module needs imported with `@_spi(Testing) import` to use this method.
+  @_spi(Internal)
+  public func withClient(
+    _ callback: @Sendable (MQTTClient) async throws -> Void
+  ) async throws {
+    try await _withClient(callback)
+  }
+
+  /// Represents connection events that clients can listen for and
+  /// react accordingly.
+  public enum Event: Sendable {
+    case connected
+    case disconnected
+    case shuttingDown
+  }
+
+  /// Represents the parameters required to publish a new value to the
+  /// MQTT broker.
+  public struct PublishRequest: Equatable, Sendable {
+
+    /// The topic to publish the new value to.
+    public let topicName: String
+
+    /// The value to publish.
+    public let payload: ByteBuffer
+
+    /// The qos of the request.
+    public let qos: MQTTQoS
+
+    /// The retain flag for the request.
+    public let retain: Bool
+
+    /// Create a new publish request.
+    ///
+    /// - Parameters:
+    ///   - topicName: The topic to publish to.
+    ///   - payload: The value to publish.
+    ///   - qos: The qos of the request.
+    ///   - retain: The retain flag of the request.
+    public init(
+      topicName: String,
+      payload: ByteBuffer,
+      qos: MQTTQoS,
+      retain: Bool
+    ) {
+      self.topicName = topicName
+      self.payload = payload
+      self.qos = qos
+      self.retain = retain
+    }
+  }
+
+}
+
+public extension MQTTManager {
+  /// Create the live manager.
+  ///
+  static func live(
+    client: MQTTClient,
+    cleanSession: Bool = false,
+    logger: Logger? = nil,
+    alwaysReconnect: Bool = true
+  ) -> Self {
+    let manager = ConnectionManager(
+      client: client,
+      logger: logger,
+      alwaysReconnect: alwaysReconnect
+    )
+    return .init(
+      connect: { try await manager.connect(cleanSession: cleanSession) },
+      connectionStream: {
+        MQTTConnectionStream(client: client, logger: logger)
+          .start()
+          .removeDuplicates()
+          .eraseToStream()
+      },
+      _listen: { topics, qos in
+        try await manager.listen(to: topics, qos: qos)
+      },
+      publish: { request in
+        let topic = request.topicName
+        guard client.isActive() else {
+          logger?.debug("Client is not active, unable to publish to topic: \(topic)")
+          return
+        }
+        logger?.trace("Begin publishing to topic: \(topic)")
+        defer { logger?.trace("Done publishing to topic: \(topic)") }
+        try await client.publish(
+          to: request.topicName,
+          payload: request.payload,
+          qos: request.qos,
+          retain: request.retain
+        )
+      },
+      shutdown: {
+        manager.shutdown()
+      },
+      _withClient: { callback in
+        try await callback(client)
+      }
+    )
+  }
+}
+
+extension MQTTManager: TestDependencyKey {
+  public static let testValue: MQTTManager = Self()
+}

Sources/MQTTManager/Internal/ConnectionManager.swift

@@ -0,0 +1,98 @@
+import Foundation
+import Logging
+import MQTTNIO
+
+actor ConnectionManager {
+  private let client: MQTTClient
+  private let logger: Logger?
+  private let name: String
+  private let shouldReconnect: Bool
+  private var hasConnected: Bool = false
+  private var listeners: [TopicListenerStream] = []
+  private var isShuttingDown = false
+
+  init(
+    client: MQTTClient,
+    logger: Logger?,
+    alwaysReconnect: Bool
+  ) {
+    var logger = logger
+    logger?[metadataKey: "instance"] = "\(Self.self)"
+    self.logger = logger
+
+    self.client = client
+    self.name = UUID().uuidString
+    self.shouldReconnect = alwaysReconnect
+  }
+
+  deinit {
+    if !isShuttingDown {
+      let message = """
+      Did not properly close the connection manager. This can lead to
+      dangling references.
+
+      Please call `shutdown` to properly close all connections and listener streams.
+      """
+      logger?.warning("\(message)")
+      self.shutdown()
+    }
+  }
+
+  private func setHasConnected() {
+    hasConnected = true
+  }
+
+  func listen(
+    to topics: [String],
+    qos: MQTTQoS
+  ) async throws -> MQTTManager.ListenStream {
+    let listener = TopicListenerStream(client: client, logger: logger, topics: topics, qos: qos)
+    listeners.append(listener)
+    await listener.start()
+    return listener.stream
+  }
+
+  func connect(
+    cleanSession: Bool
+  ) async throws {
+    guard !hasConnected else { return }
+    do {
+      try await client.connect(cleanSession: cleanSession)
+      setHasConnected()
+
+      client.addCloseListener(named: name) { [weak self] _ in
+        guard let `self` else { return }
+        self.logger?.debug("Connection closed.")
+        if self.shouldReconnect {
+          self.logger?.debug("Reconnecting...")
+          Task {
+            try await self.connect(cleanSession: cleanSession)
+          }
+        }
+      }
+
+      client.addShutdownListener(named: name) { _ in
+        self.shutdown()
+      }
+
+    } catch {
+      logger?.trace("Failed to connect: \(error)")
+      throw error
+    }
+  }
+
+  private func shutdownListeners() {
+    _ = listeners.map { $0.shutdown() }
+    listeners = []
+    isShuttingDown = true
+  }
+
+  nonisolated func shutdown(withLogging: Bool = true) {
+    if withLogging {
+      logger?.trace("Shutting down connection.")
+    }
+    client.removeCloseListener(named: name)
+    client.removeShutdownListener(named: name)
+    Task { await shutdownListeners() }
+  }
+}

Sources/MQTTManager/Internal/ConnectionStream.swift

@@ -0,0 +1,74 @@
+import Foundation
+import Logging
+import MQTTNIO
+
+@_spi(Internal)
+public actor MQTTConnectionStream: Sendable {
+
+  public typealias Element = MQTTManager.Event
+
+  private let client: MQTTClient
+  private let continuation: AsyncStream<Element>.Continuation
+  private let logger: Logger?
+  nonisolated let name: String
+  private let stream: AsyncStream<Element>
+  private var isShuttingDown = false
+
+  public init(client: MQTTClient, logger: Logger?) {
+    var logger = logger
+    logger?[metadataKey: "type"] = "\(Self.self)"
+    self.logger = logger
+
+    let (stream, continuation) = AsyncStream<Element>.makeStream()
+    self.client = client
+    self.continuation = continuation
+    self.name = UUID().uuidString
+    self.stream = stream
+  }
+
+  deinit { stop() }
+
+  public nonisolated func start() -> AsyncStream<Element> {
+    // Check if the client is active and yield the initial result.
+    continuation.yield(client.isActive() ? .connected : .disconnected)
+
+    // Continually check if the client is active.
+    let task = Task {
+      let isShuttingDown = await self.isShuttingDown
+      while !Task.isCancelled, !isShuttingDown {
+        try await Task.sleep(for: .milliseconds(100))
+        continuation.yield(client.isActive() ? .connected : .disconnected)
+      }
+    }
+
+    // Register listener on the client for when the connection
+    // closes.
+    client.addCloseListener(named: name) { _ in
+      self.logger?.trace("Client has disconnected.")
+      self.continuation.yield(.disconnected)
+    }
+
+    // Register listener on the client for when the client
+    // is shutdown.
+    client.addShutdownListener(named: name) { _ in
+      self.logger?.trace("Client is shutting down, ending connection stream: \(self.name)")
+      self.continuation.yield(.shuttingDown)
+      Task { await self.setIsShuttingDown() }
+      task.cancel()
+      self.stop()
+    }
+
+    return stream
+  }
+
+  private func setIsShuttingDown() {
+    isShuttingDown = true
+  }
+
+  public nonisolated func stop() {
+    client.removeCloseListener(named: name)
+    client.removeShutdownListener(named: name)
+    continuation.finish()
+  }
+
+}

Sources/MQTTManager/Internal/TopicListenerStream.swift

@@ -0,0 +1,177 @@
+import Foundation
+import Logging
+import MQTTNIO
+
+actor TopicListenerStream {
+
+  typealias Stream = MQTTManager.ListenStream
+
+  private let client: MQTTClient
+  private let configuration: Configuration
+  private let continuation: Stream.Continuation
+  private let logger: Logger?
+  private let name: String
+  let stream: Stream
+  private var shuttingDown: Bool = false
+  private var onShutdownHandler: (@Sendable () -> Void)?
+
+  init(
+    client: MQTTClient,
+    logger: Logger?,
+    topics: [String],
+    qos: MQTTQoS
+  ) {
+    // Setup the logger so we can more easily decipher log messages.
+    var logger = logger
+    logger?[metadataKey: "type"] = "\(Self.self)"
+    self.logger = logger
+
+    let (stream, continuation) = Stream.makeStream()
+    self.client = client
+    self.configuration = .init(qos: qos, topics: topics)
+    self.continuation = continuation
+    self.name = UUID().uuidString
+    self.stream = stream
+  }
+
+  struct Configuration: Sendable {
+    let qos: MQTTQoS
+    let topics: [String]
+  }
+
+  deinit {
+    if !shuttingDown {
+      let message = """
+      Shutdown was not called on topic listener. This could lead to potential errors or
+      the stream never ending.
+
+      Please ensure that you call shutdown on the listener.
+      """
+      client.logger.warning("\(message)")
+      continuation.finish()
+    }
+    client.removePublishListener(named: name)
+    client.removeShutdownListener(named: name)
+  }
+
+  private func subscribe() async throws {
+    guard !shuttingDown else { return }
+    logger?.debug("Begin subscribing to topics.")
+    do {
+      _ = try await client.subscribe(to: configuration.topics.map {
+        MQTTSubscribeInfo(topicFilter: $0, qos: configuration.qos)
+      })
+    } catch {
+      logger?.error("Received error while subscribing to topics: \(configuration.topics)")
+      throw TopicListenerError.failedToSubscribe
+    }
+    logger?.debug("Done subscribing to topics.")
+  }
+
+  public func start() {
+    logger?.trace("Starting listener for topics: \(configuration.topics)")
+
+    let stream = MQTTConnectionStream(client: client, logger: logger)
+      .start()
+      .removeDuplicates()
+      .eraseToStream()
+
+    let task = Task {
+      // Listen for connection events to restablish the stream upon a
+      // client becoming disconnected / reconnected, and properly shutdown
+      // the stream on the client being shutdown.
+      for await event in stream {
+        logger?.trace("Received event: \(event)")
+        switch event {
+        case .shuttingDown:
+          shutdown()
+        case .disconnected:
+          try await Task.sleep(for: .milliseconds(100))
+        case .connected:
+          try await subscribe()
+          client.addPublishListener(named: name) { result in
+            switch result {
+            case let .failure(error):
+              self.logger?.error("Received error while listening: \(error)")
+            case let .success(publishInfo):
+              if self.configuration.topics.contains(publishInfo.topicName) {
+                self.logger?.debug("Recieved new value for topic: \(publishInfo.topicName)")
+                self.continuation.yield(publishInfo)
+              }
+            }
+          }
+        }
+      }
+    }
+
+    onShutdownHandler = { task.cancel() }
+  }
+
+  // TODO: remove.
+  func listen(
+    _ topics: [String],
+    _ qos: MQTTQoS = .atLeastOnce
+  ) async throws -> Stream {
+    var sleepTimes = 0
+
+    while !client.isActive() {
+      guard sleepTimes < 10 else {
+        throw TopicListenerError.connectionTimeout
+      }
+      try? await Task.sleep(for: .milliseconds(100))
+      sleepTimes += 1
+    }
+
+    client.logger.trace("Client is active, begin subscribing to topics.")
+
+    try await subscribe()
+
+    client.logger.trace("Done subscribing, begin listening to topics.")
+
+    client.addPublishListener(named: name) { result in
+      switch result {
+      case let .failure(error):
+        self.logger?.error("Received error while listening: \(error)")
+      case let .success(publishInfo):
+        if topics.contains(publishInfo.topicName) {
+          self.logger?.debug("Recieved new value for topic: \(publishInfo.topicName)")
+          self.continuation.yield(publishInfo)
+        }
+      }
+    }
+
+    return stream
+  }
+
+  private func setIsShuttingDown() {
+    shuttingDown = true
+    onShutdownHandler = nil
+  }
+
+  public nonisolated func shutdown() {
+    client.logger.trace("Closing topic listener...")
+    continuation.finish()
+    client.removePublishListener(named: name)
+    client.removeShutdownListener(named: name)
+    Task {
+      await onShutdownHandler?()
+      await self.setIsShuttingDown()
+    }
+  }
+
+}
+
+// MARK: - Errors
+
+public enum TopicListenerError: Error {
+  case connectionTimeout
+  case failedToSubscribe
+}
+
+public struct MQTTListenResultError: Error {
+  let underlyingError: any Error
+
+  init(_ underlyingError: any Error) {
+    self.underlyingError = underlyingError
+  }
+}