OnClosed callback for producer (#293)

* Handle on_closed callback
* Implement timeout. Unconfirmed messages on OnClosed

Author: allevo

examples/ha_producer.rs

@@ -0,0 +1,151 @@
+use core::panic;
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+use std::time::Duration;
+
+use rabbitmq_stream_client::error::{ProducerPublishError, StreamCreateError};
+use rabbitmq_stream_client::types::{ByteCapacity, Message, ResponseCode};
+use rabbitmq_stream_client::Environment;
+use rabbitmq_stream_client::{
+    ConfirmationStatus, NoDedup, OnClosed, Producer, RabbitMQStreamResult,
+};
+use tokio::sync::Notify;
+use tokio::sync::RwLock;
+use tokio::time::sleep;
+use tracing::info;
+
+struct MyHAProducerInner {
+    environment: Environment,
+    stream: String,
+    producer: RwLock<Producer<NoDedup>>,
+    notify: Notify,
+    is_opened: AtomicBool,
+}
+
+#[derive(Clone)]
+struct MyHAProducer(Arc<MyHAProducerInner>);
+
+#[async_trait::async_trait]
+impl OnClosed for MyHAProducer {
+    async fn on_closed(&self, unconfirmed: Vec<Message>) {
+        info!("Producer is closed. Creating new one");
+
+        self.0
+            .is_opened
+            .store(false, std::sync::atomic::Ordering::SeqCst);
+
+        let mut producer = self.0.producer.write().await;
+
+        let new_producer = self
+            .0
+            .environment
+            .producer()
+            .build(&self.0.stream)
+            .await
+            .unwrap();
+
+        new_producer.set_on_closed(Box::new(self.clone())).await;
+
+        if !unconfirmed.is_empty() {
+            info!("Resending {} unconfirmed messages.", unconfirmed.len());
+            if let Err(e) = producer.batch_send_with_confirm(unconfirmed).await {
+                eprintln!("Error resending unconfirmed messages: {:?}", e);
+            }
+        }
+
+        *producer = new_producer;
+
+        self.0
+            .is_opened
+            .store(true, std::sync::atomic::Ordering::SeqCst);
+        self.0.notify.notify_waiters();
+    }
+}
+
+impl MyHAProducer {
+    async fn new(environment: Environment, stream: &str) -> RabbitMQStreamResult<Self> {
+        ensure_stream_exists(&environment, stream).await?;
+
+        let producer = environment.producer().build(stream).await.unwrap();
+
+        let inner = MyHAProducerInner {
+            environment,
+            stream: stream.to_string(),
+            producer: RwLock::new(producer),
+            notify: Notify::new(),
+            is_opened: AtomicBool::new(true),
+        };
+        let s = Self(Arc::new(inner));
+
+        let p = s.0.producer.write().await;
+        p.set_on_closed(Box::new(s.clone())).await;
+        drop(p);
+
+        Ok(s)
+    }
+
+    async fn send_with_confirm(
+        &self,
+        message: Message,
+    ) -> Result<ConfirmationStatus, ProducerPublishError> {
+        if !self.0.is_opened.load(std::sync::atomic::Ordering::SeqCst) {
+            self.0.notify.notified().await;
+        }
+
+        let producer = self.0.producer.read().await;
+        let err = producer.send_with_confirm(message.clone()).await;
+
+        match err {
+            Ok(s) => Ok(s),
+            Err(e) => match e {
+                ProducerPublishError::Timeout | ProducerPublishError::Closed => {
+                    Box::pin(self.send_with_confirm(message)).await
+                }
+                _ => return Err(e),
+            },
+        }
+    }
+}
+
+async fn ensure_stream_exists(environment: &Environment, stream: &str) -> RabbitMQStreamResult<()> {
+    let create_response = environment
+        .stream_creator()
+        .max_length(ByteCapacity::GB(5))
+        .create(stream)
+        .await;
+
+    if let Err(e) = create_response {
+        if let StreamCreateError::Create { stream, status } = e {
+            match status {
+                // we can ignore this error because the stream already exists
+                ResponseCode::StreamAlreadyExists => {}
+                err => {
+                    panic!("Error creating stream: {:?} {:?}", stream, err);
+                }
+            }
+        }
+    }
+
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let environment = Environment::builder().build().await?;
+    let stream = "hello-rust-stream";
+
+    let producer = MyHAProducer::new(environment, stream).await?;
+
+    let number_of_messages = 1000000;
+    for i in 0..number_of_messages {
+        let msg = Message::builder()
+            .body(format!("stream message_{}", i))
+            .build();
+        producer.send_with_confirm(msg).await?;
+        sleep(Duration::from_millis(100)).await;
+    }
+
+    Ok(())
+}

src/client/dispatcher.rs

@@ -19,22 +19,6 @@ use super::{channel::ChannelReceiver, handler::MessageHandler};
 #[derive(Clone)]
 pub(crate) struct Dispatcher<T>(DispatcherState<T>);
 
-pub(crate) struct DispatcherState<T> {
-    requests: Arc<RequestsMap>,
-    correlation_id: Arc<AtomicU32>,
-    handler: Arc<RwLock<Option<T>>>,
-}
-
-impl<T> Clone for DispatcherState<T> {
-    fn clone(&self) -> Self {
-        DispatcherState {
-            requests: self.requests.clone(),
-            correlation_id: self.correlation_id.clone(),
-            handler: self.handler.clone(),
-        }
-    }
-}
-
 struct RequestsMap {
     requests: DashMap<u32, Sender<Response>>,
     closed: AtomicBool,
@@ -126,6 +110,22 @@ where
     }
 }
 
+pub(crate) struct DispatcherState<T> {
+    requests: Arc<RequestsMap>,
+    correlation_id: Arc<AtomicU32>,
+    handler: Arc<RwLock<Option<T>>>,
+}
+
+impl<T> Clone for DispatcherState<T> {
+    fn clone(&self) -> Self {
+        DispatcherState {
+            requests: self.requests.clone(),
+            correlation_id: self.correlation_id.clone(),
+            handler: self.handler.clone(),
+        }
+    }
+}
+
 impl<T> DispatcherState<T>
 where
     T: MessageHandler,

src/client/message.rs

@@ -20,7 +20,7 @@ impl BaseMessage for Message {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct ClientMessage {
     publishing_id: u64,
     message: Message,
@@ -39,6 +39,10 @@ impl ClientMessage {
     pub fn filter_value_extract(&mut self, filter_value_extractor: impl Fn(&Message) -> String) {
         self.filter_value = Some(filter_value_extractor(&self.message));
     }
+
+    pub fn into_message(self) -> Message {
+        self.message
+    }
 }
 
 impl BaseMessage for ClientMessage {