Make dropped_subscriptions to atomic

nats-io · May 19, 2022 · f6835a1 · f6835a1
1 parent 34f7271
commit f6835a1
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 35 deletions.
diff --git a/async-nats/src/lib.rs b/async-nats/src/lib.rs
@@ -239,6 +239,7 @@ pub enum Command {
         subject: String,
         queue_group: Option<String>,
         sender: mpsc::Sender<Message>,
+        dropped_messages: Arc<AtomicU64>,
     },
     Unsubscribe {
         sid: u64,
@@ -250,10 +251,6 @@ pub enum Command {
     },
     TryFlush,
     Connect(ConnectInfo),
-    GetDroppedMessages {
-        response: tokio::sync::oneshot::Sender<u64>,
-        sid: u64,
-    },
 }
 
 /// `ClientOp` represents all actions of `Client`.
@@ -422,7 +419,7 @@ struct Subscription {
     sender: mpsc::Sender<Message>,
     queue_group: Option<String>,
     delivered: u64,
-    dropped: u64,
+    dropped_messages: Arc<AtomicU64>,
     max: Option<u64>,
 }
 
@@ -536,7 +533,9 @@ impl ConnectionHandler {
                             }
                         }
                         Err(mpsc::error::TrySendError::Full(_)) => {
-                            subscription.dropped += 1;
+                            subscription
+                                .dropped_messages
+                                .fetch_add(1, Ordering::Relaxed);
                             self.events.send(ServerEvent::SlowConsumer(sid)).await.ok();
                         }
                         Err(mpsc::error::TrySendError::Closed(_)) => {
@@ -625,11 +624,12 @@ impl ConnectionHandler {
                 subject,
                 queue_group,
                 sender,
+                dropped_messages,
             } => {
                 let subscription = Subscription {
                     sender,
                     delivered: 0,
-                    dropped: 0,
+                    dropped_messages,
                     max: None,
                     subject: subject.to_owned(),
                     queue_group: queue_group.to_owned(),
@@ -678,21 +678,6 @@ impl ConnectionHandler {
                     self.handle_disconnect().await?;
                 }
             }
-            Command::GetDroppedMessages { response, sid } => {
-                response
-                    .send(
-                        self.subscriptions
-                            .get(&sid)
-                            .ok_or_else(|| {
-                                std::io::Error::new(
-                                    ErrorKind::NotFound,
-                                    format!("could not find subscription {}", sid),
-                                )
-                            })?
-                            .dropped,
-                    )
-                    .ok();
-            }
         }
 
         Ok(())
@@ -877,6 +862,8 @@ impl Client {
         subject: String,
         queue_group: Option<String>,
     ) -> Result<Subscriber, io::Error> {
+        let dropped_messages = Arc::new(AtomicU64::new(0));
+        // dropped_messages.fetch_add
         let sid = self.next_subscription_id.fetch_add(1, Ordering::Relaxed);
         let (sender, receiver) = mpsc::channel(self.subscription_capacity);
 
@@ -886,11 +873,17 @@ impl Client {
                 subject,
                 queue_group,
                 sender,
+                dropped_messages: dropped_messages.clone(),
             })
             .await
             .unwrap();
 
-        Ok(Subscriber::new(sid, self.sender.clone(), receiver))
+        Ok(Subscriber::new(
+            sid,
+            self.sender.clone(),
+            receiver,
+            dropped_messages,
+        ))
     }
 
     pub async fn flush(&self) -> Result<(), Error> {
@@ -1093,18 +1086,21 @@ pub struct Subscriber {
     sid: u64,
     receiver: mpsc::Receiver<Message>,
     sender: mpsc::Sender<Command>,
+    dropped_messages: Arc<AtomicU64>,
 }
 
 impl Subscriber {
     fn new(
         sid: u64,
         sender: mpsc::Sender<Command>,
         receiver: mpsc::Receiver<Message>,
+        dropped_messages: Arc<AtomicU64>,
     ) -> Subscriber {
         Subscriber {
             sid,
             sender,
             receiver,
+            dropped_messages,
         }
     }
 
@@ -1183,17 +1179,8 @@ impl Subscriber {
     ///
     /// # Ok(())
     /// # }
-    pub async fn dropped(&mut self) -> io::Result<u64> {
-        let (tx, rx) = tokio::sync::oneshot::channel();
-        self.sender
-            .send(Command::GetDroppedMessages {
-                response: tx,
-                sid: self.sid,
-            })
-            .await
-            .map_err(|err| io::Error::new(ErrorKind::Other, err))?;
-        rx.await
-            .map_err(|err| io::Error::new(ErrorKind::Other, err))
+    pub async fn dropped(&mut self) -> u64 {
+        self.dropped_messages.load(Ordering::Relaxed)
     }
 }
 

diff --git a/async-nats/tests/client_tests.rs b/async-nats/tests/client_tests.rs
@@ -432,7 +432,7 @@ mod client {
             .unwrap();
         client.flush().await.unwrap();
         tokio::time::sleep(Duration::from_secs(1)).await;
-        let dropped = sub.dropped().await.unwrap();
+        let dropped = sub.dropped().await;
         println!("dropped: {}", dropped);
         assert_eq!(dropped, 1);