Implement super_stream (#232)

* WIP: implementing create and delete superstream commands

* implementing route and partition commands

* refactoring tests

* super_stream implementation

* superstream producer implementation

* super_stream_producer first basic version

* finalizing superstream_producer and tests

* better implementation/refactoring

* making Messages as references to allow borrowing

* implementing super_stream consumer

* some refactoring

* fixing super_stream_consumer test

* fix clippy issue

* adding, stream method to consumer delivery struct and adding examples

* fixing offset in SuperstreamConsumer

* super_stream_consumer new approach

* implementing close()

* implement filtering on super_stream

* adding super_stream_filtering tests

* error handling improvement in super_stream send

* improve super_stream send example

* add information to the example

Signed-off-by: Gabriele Santomaggio <G.santomaggio@gmail.com>

* small performance improvement in HashRoutingMurmurStrategy

* add information to the example

Signed-off-by: Gabriele Santomaggio <G.santomaggio@gmail.com>

* close super stream consumer

Signed-off-by: Gabriele Santomaggio <G.santomaggio@gmail.com>

* fix bug in super_stream_producer client created twice
super_stream_consumer super_stream locator not closed

* merging

* updating README.md

---------

Signed-off-by: Gabriele Santomaggio <G.santomaggio@gmail.com>
Co-authored-by: Gabriele Santomaggio <G.santomaggio@gmail.com>

Author: DanielePalaia

Cargo.toml

@@ -36,6 +36,7 @@ thiserror = "1.0"
 async-trait = "0.1.51"
 rand = "0.8"
 dashmap = "5.3.4"
+murmur3 = "0.5.2"
 
 [dev-dependencies]
 tracing-subscriber = "0.3.1"

README.md

@@ -136,6 +136,22 @@ sleep(Duration::from_secs(10)).await;
 handle.close().await?;
 ```
 
+### Superstreams
+
+The client supports the superstream functionality.
+
+A super stream is a logical stream made of individual, regular streams. It is a way to scale out publishing and consuming with RabbitMQ Streams: a large logical stream is divided into partition streams, splitting up the storage and the traffic on several cluster nodes.
+
+See the [blog post](https://blog.rabbitmq.com/posts/2022/07/rabbitmq-3-11-feature-preview-super-streams/) for more info.
+
+You can use SuperStreamProducer and SuperStreamConsumer classes which internally uses producers and consumers to operate on the componsing streams.
+
+Have a look to the examples to see on how to work with super streams.
+
+See the [Super Stream Producer Example:](https://github.com/rabbitmq/rabbitmq-stream-rust-client/examples/send_super_stream.rs)
+
+See the [Super Stream Consumer Example:](https://github.com/rabbitmq/rabbitmq-stream-rust-client/examples/receive_super_stream.rs)
+
 ### Development
 
 #### Compiling

examples/receive_super_stream.rs

@@ -0,0 +1,63 @@
+use futures::StreamExt;
+use rabbitmq_stream_client::error::StreamCreateError;
+use rabbitmq_stream_client::types::{
+    ByteCapacity, OffsetSpecification, ResponseCode, SuperStreamConsumer,
+};
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    use rabbitmq_stream_client::Environment;
+    let environment = Environment::builder().build().await?;
+    let message_count = 100_000;
+    let super_stream = "hello-rust-super-stream";
+
+    let create_response = environment
+        .stream_creator()
+        .max_length(ByteCapacity::GB(5))
+        .create_super_stream(super_stream, 3, None)
+        .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 => {
+                    println!("Error creating stream: {:?} {:?}", stream, err);
+                }
+            }
+        }
+    }
+    println!(
+        "Super stream consumer example, consuming messages from the super stream {}",
+        super_stream
+    );
+    let mut super_stream_consumer: SuperStreamConsumer = environment
+        .super_stream_consumer()
+        .offset(OffsetSpecification::First)
+        .build(super_stream)
+        .await
+        .unwrap();
+
+    for _ in 0..message_count {
+        let delivery = super_stream_consumer.next().await.unwrap();
+        {
+            let delivery = delivery.unwrap();
+            println!(
+                "Got message: {:#?} from stream: {} with offset: {}",
+                delivery
+                    .message()
+                    .data()
+                    .map(|data| String::from_utf8(data.to_vec()).unwrap())
+                    .unwrap(),
+                delivery.stream(),
+                delivery.offset()
+            );
+        }
+    }
+
+    println!("Stopping super stream consumer...");
+    let _ = super_stream_consumer.handle().close().await;
+    println!("Super stream consumer stopped");
+    Ok(())
+}

examples/send_super_stream.rs

@@ -0,0 +1,110 @@
+use rabbitmq_stream_client::error::StreamCreateError;
+use rabbitmq_stream_client::types::{
+    ByteCapacity, HashRoutingMurmurStrategy, Message, ResponseCode, RoutingStrategy,
+};
+use std::convert::TryInto;
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::sync::Arc;
+use tokio::sync::Notify;
+
+fn hash_strategy_value_extractor(message: &Message) -> String {
+    message
+        .application_properties()
+        .unwrap()
+        .get("id")
+        .unwrap()
+        .clone()
+        .try_into()
+        .unwrap()
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    use rabbitmq_stream_client::Environment;
+    let environment = Environment::builder().build().await?;
+    let message_count = 100_000;
+    let super_stream = "hello-rust-super-stream";
+    let confirmed_messages = Arc::new(AtomicU32::new(0));
+    let notify_on_send = Arc::new(Notify::new());
+    let _ = environment
+        .stream_creator()
+        .max_length(ByteCapacity::GB(5))
+        .create_super_stream(super_stream, 3, None)
+        .await;
+
+    let delete_stream = environment.delete_super_stream(super_stream).await;
+
+    match delete_stream {
+        Ok(_) => {
+            println!("Successfully deleted super stream {}", super_stream);
+        }
+        Err(err) => {
+            println!(
+                "Failed to delete super stream {}. error {}",
+                super_stream, err
+            );
+        }
+    }
+
+    let create_response = environment
+        .stream_creator()
+        .max_length(ByteCapacity::GB(5))
+        .create_super_stream(super_stream, 3, None)
+        .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 => {
+                    println!("Error creating stream: {:?} {:?}", stream, err);
+                }
+            }
+        }
+    }
+    println!(
+        "Super stream example. Sending {} messages to the super stream: {}",
+        message_count, super_stream
+    );
+    let mut super_stream_producer = environment
+        .super_stream_producer(RoutingStrategy::HashRoutingStrategy(
+            HashRoutingMurmurStrategy {
+                routing_extractor: &hash_strategy_value_extractor,
+            },
+        ))
+        .build(super_stream)
+        .await
+        .unwrap();
+
+    for i in 0..message_count {
+        let counter = confirmed_messages.clone();
+        let notifier = notify_on_send.clone();
+        let msg = Message::builder()
+            .body(format!("super stream message_{}", i))
+            .application_properties()
+            .insert("id", i.to_string())
+            .message_builder()
+            .build();
+        super_stream_producer
+            .send(msg, move |_| {
+                let inner_counter = counter.clone();
+                let inner_notifier = notifier.clone();
+                async move {
+                    if inner_counter.fetch_add(1, Ordering::Relaxed) == message_count - 1 {
+                        inner_notifier.notify_one();
+                    }
+                }
+            })
+            .await
+            .unwrap();
+    }
+
+    notify_on_send.notified().await;
+    println!(
+        "Successfully sent {} messages to the super stream {}",
+        message_count, super_stream
+    );
+    let _ = super_stream_producer.close().await;
+    Ok(())
+}

src/consumer.rs

@@ -263,6 +263,10 @@ 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?;
@@ -327,6 +331,7 @@ impl MessageHandler for ConsumerMessageHandler {
                             .0
                             .sender
                             .send(Ok(Delivery {
+                                stream: self.0.stream.clone(),
                                 subscription_id: self.0.subscription_id,
                                 message,
                                 offset,
@@ -355,6 +360,7 @@ impl MessageHandler for ConsumerMessageHandler {
 /// Envelope from incoming message
 #[derive(Debug)]
 pub struct Delivery {
+    stream: String,
     subscription_id: u8,
     message: Message,
     offset: u64,
@@ -366,6 +372,11 @@ impl Delivery {
         self.subscription_id
     }
 
+    /// Get a reference to the delivery's stream name.
+    pub fn stream(&self) -> &String {
+        &self.stream
+    }
+
     /// Get a reference to the delivery's message.
     pub fn message(&self) -> &Message {
         &self.message

src/environment.rs

@@ -11,6 +11,9 @@ use crate::{
     error::StreamDeleteError,
     producer::ProducerBuilder,
     stream_creator::StreamCreator,
+    superstream::RoutingStrategy,
+    superstream_consumer::SuperStreamConsumerBuilder,
+    superstream_producer::SuperStreamProducerBuilder,
     RabbitMQStreamResult,
 };
 
@@ -49,6 +52,18 @@ impl Environment {
         }
     }
 
+    pub fn super_stream_producer(
+        &self,
+        routing_strategy: RoutingStrategy,
+    ) -> SuperStreamProducerBuilder<NoDedup> {
+        SuperStreamProducerBuilder {
+            environment: self.clone(),
+            data: PhantomData,
+            filter_value_extractor: None,
+            route_strategy: routing_strategy,
+        }
+    }
+
     /// Returns a builder for creating a consumer
     pub fn consumer(&self) -> ConsumerBuilder {
         ConsumerBuilder {
@@ -58,6 +73,15 @@ impl Environment {
             filter_configuration: None,
         }
     }
+
+    pub fn super_stream_consumer(&self) -> SuperStreamConsumerBuilder {
+        SuperStreamConsumerBuilder {
+            environment: self.clone(),
+            offset_specification: OffsetSpecification::Next,
+            filter_configuration: None,
+        }
+    }
+
     pub(crate) async fn create_client(&self) -> RabbitMQStreamResult<Client> {
         Client::connect(self.options.client_options.clone()).await
     }

src/error.rs

@@ -110,6 +110,26 @@ pub enum ProducerPublishError {
     #[error(transparent)]
     Client(#[from] ClientError),
 }
+
+#[derive(Error, Debug)]
+pub enum SuperStreamProducerPublishError {
+    #[error("Failed to send message to stream")]
+    ProducerPublishError(),
+    #[error("Failed to create a producer")]
+    ProducerCreateError(),
+}
+
+impl From<ProducerPublishError> for SuperStreamProducerPublishError {
+    fn from(_err: ProducerPublishError) -> Self {
+        SuperStreamProducerPublishError::ProducerPublishError()
+    }
+}
+impl From<ProducerCreateError> for SuperStreamProducerPublishError {
+    fn from(_err: ProducerCreateError) -> Self {
+        SuperStreamProducerPublishError::ProducerCreateError()
+    }
+}
+
 #[derive(Error, Debug)]
 pub enum ProducerCloseError {
     #[error("Failed to close producer for stream {stream} status {status:?}")]

src/lib.rs

@@ -79,6 +79,9 @@ pub mod error;
 mod offset_specification;
 mod producer;
 mod stream_creator;
+mod superstream;
+mod superstream_consumer;
+mod superstream_producer;
 
 pub type RabbitMQStreamResult<T> = Result<T, error::ClientError>;
 
@@ -94,6 +97,10 @@ pub mod types {
     pub use crate::consumer::Delivery;
     pub use crate::offset_specification::OffsetSpecification;
     pub use crate::stream_creator::StreamCreator;
+    pub use crate::superstream::HashRoutingMurmurStrategy;
+    pub use crate::superstream::RoutingKeyRoutingStrategy;
+    pub use crate::superstream::RoutingStrategy;
+    pub use crate::superstream_consumer::SuperStreamConsumer;
     pub use rabbitmq_stream_protocol::message::Message;
     pub use rabbitmq_stream_protocol::{Response, ResponseCode, ResponseKind};
 

src/producer.rs

@@ -245,6 +245,14 @@ impl<T> ProducerBuilder<T> {
         self.filter_value_extractor = Some(f);
         self
     }
+
+    pub fn filter_value_extractor_arc(
+        mut self,
+        filter_value_extractor: Option<FilterValueExtractor>,
+    ) -> Self {
+        self.filter_value_extractor = filter_value_extractor;
+        self
+    }
 }
 
 pub struct MessageAccumulator {