Add queue subscribe

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

Author: piotrpio

Sources/Nats/NatsClient/NatsClient.swift

@@ -225,8 +225,7 @@ extension NatsClient {
             operation: ClientOp.publish((subject, inbox, payload, headers)))
 
         return try await withThrowingTaskGroup(
-            of: NatsMessage?.self,
-            body: { group in
+            of: NatsMessage?.self) { group in
                 group.addTask {
                     do {
                         return try await sub.makeAsyncIterator().next()
@@ -258,7 +257,7 @@ extension NatsClient {
 
                 // this should not be reachable
                 throw NatsError.ClientError.internalError("error waiting for response")
-            })
+            }
     }
 
     /// Flushes the internal buffer ensuring that all messages are sent.
@@ -277,22 +276,26 @@ extension NatsClient {
 
     /// Subscribes to a subject to receive messages.
     ///
-    /// - Parameter subject is a subject the client want's to subscribe to.
+    /// - Parameters:
+    ///  - subject:a subject the client want's to subscribe to.
+    ///  - queue: optional queue group name.
     ///
     /// - Returns a ``NatsSubscription`` allowing iteration over incoming messages.
     ///
     /// > **Throws:**
     /// >  - ``NatsError/ClientError/connectionClosed`` if the conneciton is closed.
     /// >  - ``NatsError/ClientError/io(_:)`` if there is an error sending the SUB request to the server.
-    public func subscribe(subject: String) async throws -> NatsSubscription {
+    /// > - ``NatsError/SubscriptionError/invalidSubject`` if the provided subject is invalid.
+    /// > - ``NatsError/SubscriptionError/invalidQueue`` if the provided queue group is invalid.
+    public func subscribe(subject: String, queue: String? = nil) async throws -> NatsSubscription {
         logger.info("subscribe to subject \(subject)")
         guard let connectionHandler = self.connectionHandler else {
             throw NatsError.ClientError.internalError("empty connection handler")
         }
         if case .closed = connectionHandler.state {
             throw NatsError.ClientError.connectionClosed
         }
-        return try await connectionHandler.subscribe(subject)
+        return try await connectionHandler.subscribe(subject, queue: queue)
     }
 
     /// Sends a PING to the server, returning the time it took for the server to respond.

Sources/Nats/NatsConnection.swift

@@ -819,11 +819,11 @@ class ConnectionHandler: ChannelInboundHandler {
         }
     }
 
-    internal func subscribe(_ subject: String) async throws -> NatsSubscription {
+    internal func subscribe(_ subject: String, queue: String? = nil) async throws -> NatsSubscription {
         let sid = self.subscriptionCounter.wrappingIncrementThenLoad(
             ordering: AtomicUpdateOrdering.relaxed)
-        try await write(operation: ClientOp.subscribe((sid, subject, nil)))
-        let sub = NatsSubscription(sid: sid, subject: subject, conn: self)
+        let sub = try NatsSubscription(sid: sid, subject: subject, queue: queue, conn: self)
+        try await write(operation: ClientOp.subscribe((sid, subject, queue)))
         self.subscriptions[sid] = sub
         return sub
     }

Sources/Nats/NatsError.swift

@@ -205,11 +205,17 @@ public enum NatsError {
     }
 
     public enum SubscriptionError: NatsErrorProtocol, Equatable {
+        case invalidSubject
+        case invalidQueue
         case permissionDenied
         case subscriptionClosed
 
         public var description: String {
             switch self {
+            case .invalidSubject:
+                return "nats: invalid subject name"
+            case .invalidQueue:
+                return "nats: invalid queue group name"
             case .permissionDenied:
                 return "nats: permission denied"
             case .subscriptionClosed:

Sources/Nats/NatsSubscription.swift

@@ -14,12 +14,12 @@
 import Foundation
 
 // TODO(pp): Implement slow consumer
-// TODO(pp): Add queue subscribe
 public class NatsSubscription: AsyncSequence {
     public typealias Element = NatsMessage
     public typealias AsyncIterator = SubscriptionIterator
 
     public let subject: String
+    public let queue: String?
     internal var max: UInt64?
     internal var delivered: UInt64 = 0
     internal let sid: UInt64
@@ -34,14 +34,21 @@ public class NatsSubscription: AsyncSequence {
 
     private static let defaultSubCapacity: UInt64 = 512 * 1024
 
-    convenience init(sid: UInt64, subject: String, conn: ConnectionHandler) {
-        self.init(
-            sid: sid, subject: subject, capacity: NatsSubscription.defaultSubCapacity, conn: conn)
+    convenience init(sid: UInt64, subject: String, queue: String?, conn: ConnectionHandler) throws {
+        try self.init(
+            sid: sid, subject: subject, queue: queue, capacity: NatsSubscription.defaultSubCapacity, conn: conn)
     }
 
-    init(sid: UInt64, subject: String, capacity: UInt64, conn: ConnectionHandler) {
+    init(sid: UInt64, subject: String, queue: String?, capacity: UInt64, conn: ConnectionHandler) throws {
+        if !NatsSubscription.validSubject(subject) {
+            throw NatsError.SubscriptionError.invalidSubject
+        }
+        if let queue, !NatsSubscription.validQueue(queue) {
+            throw NatsError.SubscriptionError.invalidQueue
+        }
         self.sid = sid
         self.subject = subject
+        self.queue = queue
         self.capacity = capacity
         self.buffer = []
         self.conn = conn
@@ -150,4 +157,26 @@ public class NatsSubscription: AsyncSequence {
         }
         return try await self.conn.unsubscribe(sub: self, max: after)
     }
+
+    // validateSubject will do a basic subject validation.
+    // Spaces are not allowed and all tokens should be > 0 in length.
+    private static func validSubject(_ subj: String) -> Bool {
+        let whitespaceCharacterSet = CharacterSet.whitespacesAndNewlines
+        if subj.rangeOfCharacter(from: whitespaceCharacterSet) != nil {
+            return false
+        }
+        let tokens = subj.split(separator: ".")
+        for token in tokens {
+            if token.isEmpty {
+                return false
+            }
+        }
+        return true
+    }
+
+    // validQueue will check a queue name for whitespaces.
+    private static func validQueue(_ queue: String) -> Bool {
+        let whitespaceCharacterSet = CharacterSet.whitespacesAndNewlines
+        return queue.rangeOfCharacter(from: whitespaceCharacterSet) == nil
+    }
 }

Tests/NatsTests/Integration/ConnectionTests.swift

@@ -314,6 +314,53 @@ class CoreNatsTests: XCTestCase {
         XCTAssertEqual(message?.payload, "msg".data(using: .utf8)!)
     }
 
+    func testQueueGroupSubscribe() async throws {
+        natsServer.start()
+        logger.logLevel = .debug
+        let client = NatsClientOptions().url(URL(string: natsServer.clientURL)!).build()
+        try await client.connect()
+
+        let sub1 = try await client.subscribe(subject: "test", queue: "queueGroup")
+        let sub2 = try await client.subscribe(subject: "test", queue: "queueGroup")
+
+        try await client.publish("msg".data(using: .utf8)!, subject: "test")
+
+        try await withThrowingTaskGroup(of: NatsMessage?.self) { group in
+            group.addTask {try await sub1.makeAsyncIterator().next()}
+            group.addTask {try await sub2.makeAsyncIterator().next()}
+            group.addTask {
+                try await Task.sleep(nanoseconds: UInt64(1_000_000_000))
+                return nil
+            }
+
+            var msgReceived = false
+            var timeoutReceived = false
+            for try await result in group {
+                if let _ = result {
+                    if msgReceived == true {
+                        XCTFail("received 2 messages")
+                        return
+                    }
+                    msgReceived = true
+                } else {
+                    if !msgReceived {
+                        XCTFail("timeout received before getting any messages")
+                        return
+                    }
+                    timeoutReceived = true
+                }
+                if msgReceived && timeoutReceived {
+                    break
+                }
+            }
+            group.cancelAll()
+            try await sub1.unsubscribe()
+            try await sub2.unsubscribe()
+            return
+        }
+    }
+
+
     func testUnsubscribe() async throws {
         natsServer.start()
         logger.logLevel = .debug