Add closing client connection and client ping

Signed-off-by: Piotr Piotrowski <piotr@synadia.com>

Author: piotrpio

Sources/NatsSwift/NatsClient/NatsClient.swift

@@ -17,6 +17,7 @@ public enum NatsState {
     case Pending
     case Connected
     case Disconnected
+    case Closed
 }
 
 public struct Auth {
@@ -67,6 +68,14 @@ extension Client {
         }
         try await connectionHandler.connect()
     }
+    
+    public func close() async throws {
+        logger.debug("close")
+        guard let connectionHandler = self.connectionHandler else {
+            throw NSError(domain: "nats_swift", code: 1, userInfo: ["message": "empty connection handler"])
+        }
+        try await connectionHandler.close()
+    }
 
     public func publish(_ payload: Data, subject: String, reply: String? = nil, headers: HeaderMap? = nil) throws {
         logger.debug("publish")

Sources/NatsSwift/NatsConnection.swift

@@ -7,13 +7,16 @@ import Foundation
 import NIO
 import NIOFoundationCompat
 import Dispatch
+import Atomics
 
 class ConnectionHandler: ChannelInboundHandler {
     let lang = "Swift"
     let version = "0.0.1"
 
     internal let allocator = ByteBufferAllocator()
     internal var inputBuffer: ByteBuffer
+
+    // Connection options
     internal var urls: [URL]
     // nanoseconds representation of TimeInterval
     internal let reconnectWait: UInt64
@@ -23,10 +26,12 @@ class ConnectionHandler: ChannelInboundHandler {
     typealias InboundIn = ByteBuffer
     internal var state: NatsState = .Pending
     internal var subscriptions: [ UInt64: Subscription ]
-    internal var subscriptionCounter = SubscriptionCounter()
+    internal var subscriptionCounter = ManagedAtomic<UInt64>(0)
     internal var serverInfo: ServerInfo?
     internal var auth: Auth?
     private var parseRemainder: Data?
+    private var pingTask: RepeatedTask?
+    private var outstandingPings = ManagedAtomic<UInt8>(0)
 
 
     func channelRead(context: ChannelHandlerContext, data: NIOAny) {
@@ -50,28 +55,29 @@ class ConnectionHandler: ChannelInboundHandler {
         }
         for op in parseResult.ops {
             if let continuation = self.serverInfoContinuation {
+                self.serverInfoContinuation = nil
                 logger.debug("server info")
                 switch op {
                 case let .Error(err):
                     continuation.resume(throwing: err)
                 case let .Info(info):
                     continuation.resume(returning: info)
                 default:
-                    continuation.resume(throwing: NSError(domain: "nats_swift", code: 1, userInfo: ["message": "unexpected operation; expected server info: \(op)"]))
+                    // ignore until we get either error or server info
+                    continue
                 }
-                self.serverInfoContinuation = nil
                 continue
             }
 
             if let continuation = self.connectionEstablishedContinuation {
+                self.connectionEstablishedContinuation = nil
                 logger.debug("conn established")
                 switch op {
                 case let .Error(err):
                     continuation.resume(throwing: err)
                 default:
                     continuation.resume()
                 }
-                self.connectionEstablishedContinuation = nil
                 continue
             }
 
@@ -85,6 +91,9 @@ class ConnectionHandler: ChannelInboundHandler {
                     logger.error("error sending pong: \(error)")
                     continue
                 }
+            case .Pong:
+                logger.debug("pong")
+                self.outstandingPings.store(0, ordering: AtomicStoreOrdering.relaxed)
             case let .Error(err):
                 logger.debug("error \(err)")
             case let .Message(msg):
@@ -165,9 +174,44 @@ class ConnectionHandler: ChannelInboundHandler {
             }
         }
         self.state = .Connected
+        guard let channel = self.channel else {
+            throw NSError(domain: "nats_swift", code: 1, userInfo: ["message": "empty channel"])
+        }
+        // Schedule the task to send a PING periodically
+        let pingInterval = TimeAmount.nanoseconds(Int64(self.pingInterval*1_000_000_000))
+        self.pingTask = channel.eventLoop.scheduleRepeatedTask(initialDelay: pingInterval, delay: pingInterval) { [weak self] task in
+             self?.sendPing()
+         }
         logger.debug("connection established")
     }
 
+    func close() async throws {
+        self.state = .Closed
+        try await disconnect()
+        try await self.group.shutdownGracefully()
+    }
+
+    func disconnect() async throws {
+        self.pingTask?.cancel()
+        try await self.channel?.close().get()
+    }
+
+    private func sendPing() {
+        let pingsOut = self.outstandingPings.wrappingIncrementThenLoad(ordering: AtomicUpdateOrdering.relaxed)
+        if pingsOut > 2 {
+            handleDisconnect()
+            return
+        }
+        let ping = ClientOp.Ping
+        do {
+            try self.write(operation: ping)
+            logger.debug("sent ping: \(pingsOut)")
+        } catch {
+            logger.error("Unable to send ping: \(error)")
+        }
+
+    }
+
     func channelActive(context: ChannelHandlerContext) {
         logger.debug("TCP channel active")
 
@@ -177,18 +221,43 @@ class ConnectionHandler: ChannelInboundHandler {
     func channelInactive(context: ChannelHandlerContext) {
         logger.debug("TCP channel inactive")
 
-        handleDisconnect()
+        if self.state == .Connected {
+            handleDisconnect()
+        }
     }
 
     func errorCaught(context: ChannelHandlerContext, error: Error) {
         // TODO(pp): implement Close() on the connection and call it here
         logger.debug("Encountered error on the channel: \(error)")
-        self.state = .Disconnected
-        handleReconnect()
+        context.close(promise: nil)
+        if self.state == .Connected {
+            handleDisconnect()
+        } else if self.state == .Disconnected {
+            handleReconnect()
+        }
     }
 
     func handleDisconnect() {
         self.state = .Disconnected
+        if let channel = self.channel {
+            let promise = channel.eventLoop.makePromise(of: Void.self)
+            Task {
+                do {
+                    try await self.disconnect()
+                    promise.succeed()
+                } catch {
+                    promise.fail(error)
+                }
+            }
+            promise.futureResult.whenComplete { result in
+                do {
+                  try result.get()
+                } catch {
+                    logger.error("Error closing connection: \(error)")
+                }
+            }
+        }
+
         handleReconnect()
     }
 
@@ -200,10 +269,12 @@ class ConnectionHandler: ChannelInboundHandler {
                     try await self.connect()
                 } catch {
                     // TODO(pp): add option to set this to exponential backoff (with jitter)
+                    logger.debug("could not reconnect: \(error)")
                     try await Task.sleep(nanoseconds: self.reconnectWait)
                     attempts += 1
                     continue
                 }
+                logger.debug("reconnected")
                 break
             }
             for (sid, sub) in self.subscriptions {
@@ -242,7 +313,7 @@ class ConnectionHandler: ChannelInboundHandler {
     }
 
     func subscribe(_ subject: String) async throws -> Subscription {
-        let sid = self.subscriptionCounter.next()
+        let sid = self.subscriptionCounter.wrappingIncrementThenLoad(ordering: AtomicUpdateOrdering.relaxed)
         try write(operation: ClientOp.Subscribe((sid, subject, nil)))
         let sub = Subscription(subject: subject)
         self.subscriptions[sid] = sub

Sources/NatsSwift/NatsSubscription.swift

@@ -84,15 +84,3 @@ public class Subscription: AsyncSequence {
         }
     }
 }
-
-internal class SubscriptionCounter {
-    private var counter: UInt64 = 0
-    private let queue = DispatchQueue(label: "io.nats.swift.subscriptionCounter")
-
-    func next() -> UInt64 {
-        queue.sync {
-            counter+=1
-            return counter
-        }
-    }
-}

Tests/NatsSwiftTests/Integration/ConnectionTests.swift

@@ -121,7 +121,7 @@ class CoreNatsTests: XCTestCase {
         }
 
         // make sure sub receives messages
-        for await msg in sub {
+        for await _ in sub {
             messagesReceived += 1
             if messagesReceived == 10 {
                 break
@@ -140,16 +140,17 @@ class CoreNatsTests: XCTestCase {
                 try client.publish(payload, subject: "foo")
             }
         }
-
-        for await msg in sub {
+        
+        for await _ in sub {
             messagesReceived += 1
-            if messagesReceived == 10 {
+            if messagesReceived == 20 {
                 break
             }
         }
 
         // Check if the total number of messages received matches the number sent
         XCTAssertEqual(20, messagesReceived, "Mismatch in the number of messages sent and received")
+        try await client.close()
     }
 
     func testUsernameAndPassword() async throws {
@@ -226,3 +227,4 @@ class CoreNatsTests: XCTestCase {
 
     }
 }
+