add queue subscriber and unit test

async-nats/examples/sub.rs

@@ -1,4 +1,3 @@
-use bytes::Bytes;
 use futures_util::StreamExt;
 use std::error::Error;
 use std::time::Instant;

async-nats/src/lib.rs

@@ -226,6 +226,7 @@ pub enum ClientOp {
     Subscribe {
         sid: u64,
         subject: String,
+        queue_group: Option<String>,
     },
     Unsubscribe {
         id: u64,
@@ -455,9 +456,18 @@ impl Connection {
                 self.stream.write_all(b"\r\n").await?;
             }
 
-            ClientOp::Subscribe { sid, subject } => {
+            ClientOp::Subscribe {
+                sid,
+                subject,
+                queue_group,
+            } => {
                 self.stream.write_all(b"SUB ").await?;
                 self.stream.write_all(subject.as_bytes()).await?;
+                if let Some(queue_group) = queue_group {
+                    self.stream
+                        .write_all(format!(" {}", queue_group).as_bytes())
+                        .await?;
+                }
                 self.stream
                     .write_all(format!(" {}\r\n", sid).as_bytes())
                     .await?;
@@ -725,15 +735,38 @@ impl Client {
         format!("_INBOX.{}", nuid::next())
     }
 
+    pub async fn queue_subscribe(
+        &mut self,
+        subject: String,
+        queue_group: String,
+    ) -> Result<Subscriber, io::Error> {
+        self._subscribe(subject, Some(queue_group)).await
+    }
+
     pub async fn subscribe(&mut self, subject: String) -> Result<Subscriber, io::Error> {
+        self._subscribe(subject, None).await
+    }
+
+    // TODO: options/questions for nats team:
+    //  - should there just be a single subscribe() function (would be breaking api against 0.11.0)
+    //  - if queue_subscribe is a separate function, how do you want to name the private function here?
+    async fn _subscribe(
+        &mut self,
+        subject: String,
+        queue_group: Option<String>,
+    ) -> Result<Subscriber, io::Error> {
         let (sender, receiver) = mpsc::channel(16);
 
         // Aiming to make this the only lock (aside from internal locks in channels).
         let mut context = self.subscription_context.lock().await;
         let sid = context.insert(Subscription { sender });
 
         self.sender
-            .send(ClientOp::Subscribe { sid, subject })
+            .send(ClientOp::Subscribe {
+                sid,
+                subject,
+                queue_group,
+            })
             .await
             .unwrap();
 
@@ -749,9 +782,9 @@ impl Client {
     }
 }
 
-/// Connets to the NATS with specified options.
+/// Connects to NATS with specified options.
 ///
-/// It is generally advised to use [ConnectOptions] instead, as it provides builder for whole
+/// It is generally advised to use [ConnectOptions] instead, as it provides a builder for whole
 /// configuration.
 ///
 /// # Examples

async-nats/tests/client_tests.rs

@@ -17,7 +17,7 @@ mod client {
 
     use super::nats_server;
     use bytes::Bytes;
-    use futures_util::StreamExt;
+    use futures_util::{future::join_all, StreamExt};
 
     #[tokio::test]
     async fn basic_pub_sub() {
@@ -46,6 +46,53 @@ mod client {
         assert_eq!(i, 10);
     }
 
+    #[tokio::test]
+    async fn queue_sub() {
+        let server = nats_server::run_basic_server();
+        const NUM_SUBSCRIBERS: usize = 3;
+        const NUM_ITEMS: usize = 20;
+
+        let mut subscribers = Vec::new();
+        let mut client = async_nats::connect(server.client_url()).await.unwrap();
+        for _i in 0..NUM_SUBSCRIBERS {
+            subscribers.push(
+                client
+                    .queue_subscribe("qfoo".into(), "group".into())
+                    .await
+                    .unwrap(),
+            );
+        }
+
+        for _ in 0..NUM_ITEMS {
+            client.publish("qfoo".into(), "data".into()).await.unwrap();
+        }
+        client.flush().await.unwrap();
+        let mut results = Vec::new();
+        for mut subscriber in subscribers.into_iter() {
+            results.push(tokio::spawn(async move {
+                let mut count = 0u32;
+                while let Ok(Some(item)) = tokio::time::timeout(
+                    tokio::time::Duration::from_millis(1000),
+                    subscriber.next(),
+                )
+                .await
+                {
+                    count += 1;
+                }
+                count
+            }));
+        }
+        let counts = join_all(results.iter_mut())
+            .await
+            .into_iter()
+            .filter_map(|n| n.ok())
+            .collect::<Vec<u32>>();
+        let total: u32 = counts.iter().sum();
+        assert_eq!(total, NUM_ITEMS as u32, "all items received");
+        let num_receivers = counts.into_iter().filter(|n| *n > 0u32).count();
+        assert!(num_receivers > 1, "should not all go to single subscriber");
+    }
+
     #[tokio::test]
     async fn cloned_client() {
         let server = nats_server::run_basic_server();