implementing close()

Author: DanielePalaia

examples/receive_super_stream.rs

@@ -3,11 +3,6 @@ use rabbitmq_stream_client::error::StreamCreateError;
 use rabbitmq_stream_client::types::{
     ByteCapacity, OffsetSpecification, ResponseCode, SuperStreamConsumer,
 };
-use std::sync::atomic::{AtomicU32, Ordering};
-use std::sync::Arc;
-use std::time::Duration;
-use tokio::task;
-use tokio::time::sleep;
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
@@ -41,34 +36,25 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .await
         .unwrap();
 
-    let received_messages = Arc::new(AtomicU32::new(0));
-
-    for mut consumer in super_stream_consumer.get_consumers().await.into_iter() {
-        let received_messages_outer = received_messages.clone();
-
-        task::spawn(async move {
-            let mut inner_received_messages = received_messages_outer.clone();
-            while let Some(delivery) = consumer.next().await {
-                let d = delivery.unwrap();
-                println!(
-                    "Got message: {:#?} from stream: {} with offset: {}",
-                    d.message()
-                        .data()
-                        .map(|data| String::from_utf8(data.to_vec()).unwrap()),
-                    d.stream(),
-                    d.offset(),
-                );
-                let value = inner_received_messages.fetch_add(1, Ordering::Relaxed);
-                if value == message_count {
-                    let handle = consumer.handle();
-                    _ = handle.close().await;
-                    break;
-                }
-            }
-        });
+    let mut received_messages = 0;
+
+    while let delivery = super_stream_consumer.next().await.unwrap() {
+        println!("inside while delivery loop");
+        let d = delivery.unwrap();
+        println!(
+            "Got message: {:#?} from stream: {} with offset: {}",
+            d.message()
+                .data()
+                .map(|data| String::from_utf8(data.to_vec()).unwrap()),
+            d.stream(),
+            d.offset()
+        );
+
+        received_messages = received_messages + 1;
+        if received_messages == 10 {
+            break;
+        }
     }
 
-    sleep(Duration::from_millis(20000)).await;
-
     Ok(())
 }

src/consumer.rs

@@ -261,11 +261,16 @@ impl Stream for Consumer {
 }
 
 /// Handler API for [`Consumer`]
+///
 pub struct ConsumerHandle(Arc<ConsumerInternal>);
 
 impl ConsumerHandle {
     /// Close the [`Consumer`] associated to this handle
     pub async fn close(self) -> Result<(), ConsumerCloseError> {
+        self.internal_close().await
+    }
+
+    pub(crate) async fn internal_close(&self) -> Result<(), ConsumerCloseError> {
         match self.0.closed.compare_exchange(false, true, SeqCst, SeqCst) {
             Ok(false) => {
                 let response = self.0.client.unsubscribe(self.0.subscription_id).await?;

src/superstream_consumer.rs

@@ -1,10 +1,14 @@
 use crate::consumer::Delivery;
-use crate::error::ConsumerDeliveryError;
+use crate::error::{ConsumerCloseError, ConsumerDeliveryError};
 use crate::superstream::DefaultSuperStreamMetadata;
-use crate::{error::ConsumerCreateError, Environment};
+use crate::{error::ConsumerCreateError, ConsumerHandle, Environment};
+use futures::task::AtomicWaker;
 use futures::{Stream, StreamExt};
 use rabbitmq_stream_protocol::commands::subscribe::OffsetSpecification;
 use std::pin::Pin;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering::{Relaxed, SeqCst};
+use std::sync::Arc;
 use std::task::{Context, Poll};
 use tokio::sync::mpsc::{channel, Receiver};
 use tokio::task;
@@ -13,11 +17,14 @@ use tokio::task;
 
 /// API for consuming RabbitMQ stream messages
 pub struct SuperStreamConsumer {
-    internal: SuperStreamConsumerInternal,
+    internal: Arc<SuperStreamConsumerInternal>,
+    receiver: Receiver<Result<Delivery, ConsumerDeliveryError>>,
 }
 
 struct SuperStreamConsumerInternal {
-    receiver: Receiver<Result<Delivery, ConsumerDeliveryError>>,
+    closed: Arc<AtomicBool>,
+    handlers: Vec<ConsumerHandle>,
+    waker: AtomicWaker,
 }
 
 /// Builder for [`Consumer`]
@@ -44,6 +51,7 @@ impl SuperStreamConsumerBuilder {
         };
         let partitions = super_stream_metadata.partitions().await;
 
+        let mut handlers = Vec::<ConsumerHandle>::new();
         for partition in partitions.into_iter() {
             let tx_cloned = tx.clone();
             let mut consumer = self
@@ -54,17 +62,24 @@ impl SuperStreamConsumerBuilder {
                 .await
                 .unwrap();
 
+            handlers.push(consumer.handle());
+
             task::spawn(async move {
                 while let Some(d) = consumer.next().await {
                     _ = tx_cloned.send(d).await;
                 }
             });
         }
 
-        let super_stream_consumer_internal = SuperStreamConsumerInternal { receiver: rx };
+        let super_stream_consumer_internal = SuperStreamConsumerInternal {
+            closed: Arc::new(AtomicBool::new(false)),
+            handlers,
+            waker: AtomicWaker::new(),
+        };
 
         Ok(SuperStreamConsumer {
-            internal: super_stream_consumer_internal,
+            internal: Arc::new(super_stream_consumer_internal),
+            receiver: rx,
         })
     }
 
@@ -78,6 +93,46 @@ impl Stream for SuperStreamConsumer {
     type Item = Result<Delivery, ConsumerDeliveryError>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        Pin::new(&mut self.internal.receiver).poll_recv(cx)
+        self.internal.waker.register(cx.waker());
+        let poll = Pin::new(&mut self.receiver).poll_recv(cx);
+        match (self.is_closed(), poll.is_ready()) {
+            (true, false) => Poll::Ready(None),
+            _ => poll,
+        }
+    }
+}
+
+impl SuperStreamConsumer {
+    /// Check if the consumer is closed
+    pub fn is_closed(&self) -> bool {
+        self.internal.is_closed()
+    }
+
+    pub fn handle(&self) -> SuperStreamConsumerHandle {
+        SuperStreamConsumerHandle(self.internal.clone())
+    }
+}
+
+impl SuperStreamConsumerInternal {
+    fn is_closed(&self) -> bool {
+        self.closed.load(Relaxed)
+    }
+}
+
+pub struct SuperStreamConsumerHandle(Arc<SuperStreamConsumerInternal>);
+
+impl SuperStreamConsumerHandle {
+    /// Close the [`Consumer`] associated to this handle
+    pub async fn close(self) -> Result<(), ConsumerCloseError> {
+        self.0.waker.wake();
+        match self.0.closed.compare_exchange(false, true, SeqCst, SeqCst) {
+            Ok(false) => {
+                for handle in &self.0.handlers {
+                    handle.internal_close().await.unwrap();
+                }
+                Ok(())
+            }
+            _ => Err(ConsumerCloseError::AlreadyClosed),
+        }
     }
 }

tests/integration/consumer_test.rs

@@ -98,6 +98,7 @@ async fn super_stream_consumer_test() {
     }
 
     let mut received_messages = 0;
+    let handle = super_stream_consumer.handle();
 
     println!("before looping");
     while let delivery = super_stream_consumer.next().await.unwrap() {
@@ -121,6 +122,7 @@ async fn super_stream_consumer_test() {
     assert!(received_messages == message_count);
 
     super_stream_producer.close().await.unwrap();
+    _ = handle.close().await;
 }
 
 #[tokio::test(flavor = "multi_thread")]