implementing super_stream consumer

Author: DanielePalaia

src/environment.rs

@@ -13,6 +13,7 @@ use crate::{
     stream_creator::StreamCreator,
     superstream::RoutingStrategy,
     superstream_producer::SuperStreamProducerBuilder,
+    superstream_consumer::SuperStreamConsumerBuilder,
     RabbitMQStreamResult,
 };
 
@@ -72,6 +73,14 @@ impl Environment {
             filter_configuration: None,
         }
     }
+
+    pub fn super_stream_consumer(&self) -> SuperStreamConsumerBuilder {
+        SuperStreamConsumerBuilder {
+            environment: self.clone(),
+            offset_specification: OffsetSpecification::Next,
+        }
+    }
+
     pub(crate) async fn create_client(&self) -> RabbitMQStreamResult<Client> {
         Client::connect(self.options.client_options.clone()).await
     }

src/lib.rs

@@ -80,6 +80,7 @@ mod offset_specification;
 mod producer;
 mod stream_creator;
 mod superstream;
+mod superstream_consumer;
 mod superstream_producer;
 
 pub type RabbitMQStreamResult<T> = Result<T, error::ClientError>;
@@ -98,6 +99,7 @@ pub mod types {
     pub use crate::stream_creator::StreamCreator;
     pub use crate::superstream::HashRoutingMurmurStrategy;
     pub use crate::superstream::RoutingKeyRoutingStrategy;
+    pub use crate::superstream_consumer::SuperStreamConsumer;
     pub use crate::superstream::RoutingStrategy;
     pub use rabbitmq_stream_protocol::message::Message;
     pub use rabbitmq_stream_protocol::{Response, ResponseCode, ResponseKind};

src/superstream.rs

@@ -4,13 +4,6 @@ use murmur3::murmur3_32;
 use rabbitmq_stream_protocol::message::Message;
 use std::io::Cursor;
 
-/*
-pub enum RouteType {
-    Hash = 0,
-    Key = 1,
-}
-*/
-
 #[derive(Clone)]
 pub struct DefaultSuperStreamMetadata {
     pub super_stream: String,

src/superstream_consumer.rs

@@ -0,0 +1,82 @@
+use std::sync::Arc;
+
+use rabbitmq_stream_protocol::{commands::subscribe::OffsetSpecification, message::Message};
+
+use crate::superstream::DefaultSuperStreamMetadata;
+use crate::{error::ConsumerCreateError, Client, Consumer, Environment};
+
+type FilterPredicate = Option<Arc<dyn Fn(&Message) -> bool + Send + Sync>>;
+
+/// API for consuming RabbitMQ stream messages
+#[derive(Clone)]
+pub struct SuperStreamConsumer {
+    pub internal: Arc<SuperStreamConsumerInternal>,
+}
+
+struct SuperStreamConsumerInternal {
+    client: Client,
+    super_stream: String,
+    offset_specification: OffsetSpecification,
+    pub consumers: Vec<Consumer>,
+}
+
+/// Builder for [`Consumer`]
+pub struct SuperStreamConsumerBuilder {
+    pub(crate) environment: Environment,
+    pub(crate) offset_specification: OffsetSpecification,
+}
+
+impl SuperStreamConsumerBuilder {
+    pub async fn build(
+        self,
+        super_stream: &str,
+    ) -> Result<SuperStreamConsumer, ConsumerCreateError> {
+        // Connect to the user specified node first, then look for a random replica to connect to instead.
+        // This is recommended for load balancing purposes.
+        let client = self.environment.create_client().await?;
+
+        let mut super_stream_metadata = DefaultSuperStreamMetadata {
+            super_stream: super_stream.to_string(),
+            client: client.clone(),
+            partitions: Vec::new(),
+            routes: Vec::new(),
+        };
+        let partitions = super_stream_metadata.partitions().await;
+        let mut consumers: Vec<Consumer> = Vec::new();
+
+        for partition in partitions.into_iter() {
+            let consumer = self
+                .environment
+                .consumer()
+                .offset(self.offset_specification.clone())
+                .build(partition.as_str())
+                .await
+                .unwrap();
+
+            consumers.push(consumer);
+        }
+
+        let super_stream_consumer_internal = Arc::new(SuperStreamConsumerInternal {
+            super_stream: super_stream.to_string(),
+            client: client.clone(),
+            offset_specification: self.offset_specification.clone(),
+            consumers,
+        });
+
+        Ok(SuperStreamConsumer {
+            internal: super_stream_consumer_internal,
+        })
+    }
+
+    pub fn offset(mut self, offset_specification: OffsetSpecification) -> Self {
+        self.offset_specification = offset_specification;
+        self
+    }
+}
+
+impl SuperStreamConsumer   {
+
+    pub async fn get_consumers(&self) -> &Vec<Consumer>   {
+        return &self.internal.consumers
+    }
+}

src/superstream_producer.rs

@@ -12,7 +12,7 @@ use std::future::Future;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
-type FilterValueExtractor = Arc<dyn Fn(&Message) -> String + 'static + Send + Sync>;
+//type FilterValueExtractor = Arc<dyn Fn(&Message) -> String + 'static + Send + Sync>;
 
 #[derive(Clone)]
 pub struct SuperStreamProducer<T>(
@@ -62,8 +62,8 @@ impl SuperStreamProducer<NoDedup> {
 
         for route in routes.into_iter() {
             if !self.1.contains_key(route.as_str()) {
-                    let producer = self.0.environment.producer().build(route.as_str()).await;
-                    self.1.insert(route.clone(), producer.unwrap());
+                let producer = self.0.environment.producer().build(route.as_str()).await;
+                self.1.insert(route.clone(), producer.unwrap());
             }
 
             let producer = self.1.get(route.as_str()).unwrap();
@@ -72,7 +72,6 @@ impl SuperStreamProducer<NoDedup> {
         Ok(())
     }
 
-
     pub async fn close(self) -> Result<(), ProducerCloseError> {
         self.0.client.close().await?;
 

tests/integration/consumer_test.rs

@@ -3,17 +3,20 @@ use std::time::Duration;
 use crate::common::TestEnvironment;
 use fake::{Fake, Faker};
 use futures::StreamExt;
+use tokio::task;
 use rabbitmq_stream_client::{
     error::{
         ClientError, ConsumerCloseError, ConsumerDeliveryError, ConsumerStoreOffsetError,
         ProducerCloseError,
     },
-    types::{Delivery, Message, OffsetSpecification},
+    types::{Delivery, Message, OffsetSpecification, SuperStreamConsumer},
     Consumer, FilterConfiguration, NoDedup, Producer,
 };
 
 use rabbitmq_stream_protocol::ResponseCode;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU32, Ordering};
+use rabbitmq_stream_client::types::{HashRoutingMurmurStrategy, RoutingStrategy};
 
 #[tokio::test(flavor = "multi_thread")]
 async fn consumer_test() {
@@ -54,6 +57,69 @@ async fn consumer_test() {
     producer.close().await.unwrap();
 }
 
+fn hash_strategy_value_extractor(message: &Message) -> String {
+    let s = String::from_utf8(Vec::from(message.data().unwrap())).expect("Found invalid UTF-8");
+    return s;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn super_stream_consumer_test() {
+    let env = TestEnvironment::create_super_stream().await;
+    let reference: String = Faker.fake();
+
+    let message_count = 10;
+    let mut super_stream_producer = env
+        .env
+        .super_stream_producer(RoutingStrategy::HashRoutingStrategy(
+            HashRoutingMurmurStrategy {
+                routing_extractor: &hash_strategy_value_extractor,
+            },
+        ))
+        .build(&env.super_stream)
+        .await
+        .unwrap();
+
+    static super_stream_consumer: SuperStreamConsumer = env
+        .env
+        .super_stream_consumer()
+        .offset(OffsetSpecification::Next)
+        .build(&env.stream)
+        .await
+        .unwrap();
+
+    for n in 0..message_count {
+        let msg = Message::builder().body(format!("message{}", n)).build();
+        let _ = super_stream_producer
+            .send(msg, |confirmation_status| async move {
+                println!("Message confirmed with status {:?}", confirmation_status);
+            })
+            .await
+            .unwrap();
+    }
+
+
+    let received_messages = Arc::new(AtomicU32::new(0));
+    let consumers =  super_stream_consumer.get_consumers().await;
+
+    let mut tasks = Vec::new();
+        for mut consumer in consumers.into_iter() {
+            let received_messages_outer = received_messages.clone();
+            tasks.push(task::spawn(async move {
+                let inner_received_messages = received_messages_outer.clone();
+                let delivery = consumer.next().await.unwrap();
+                let _ = String::from_utf8(delivery.unwrap().message().data().unwrap().to_vec()).unwrap();
+                inner_received_messages.fetch_add(1, Ordering::Relaxed);
+
+            }));
+        }
+
+    futures::future::join_all(tasks).await;
+
+    assert!(received_messages.fetch_add(1, Ordering::Relaxed) == message_count);
+    //consumer.handle().close().await.unwrap();
+    super_stream_producer.close().await.unwrap();
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn consumer_test_offset_specification_offset() {
     let env = TestEnvironment::create().await;

tests/integration/producer_test.rs

@@ -403,7 +403,7 @@ fn hash_strategy_value_extractor(message: &Message) -> String {
 #[tokio::test(flavor = "multi_thread")]
 async fn key_super_steam_producer_test() {
     let env = TestEnvironment::create_super_stream().await;
-    let mut confirmed_messages = Arc::new(AtomicU32::new(0));
+    let confirmed_messages = Arc::new(AtomicU32::new(0));
     let notify_on_send = Arc::new(Notify::new());
     let message_count = 100;
 
@@ -443,7 +443,7 @@ async fn key_super_steam_producer_test() {
 #[tokio::test(flavor = "multi_thread")]
 async fn hash_super_steam_producer_test() {
     let env = TestEnvironment::create_super_stream().await;
-    let mut confirmed_messages = Arc::new(AtomicU32::new(0));
+    let confirmed_messages = Arc::new(AtomicU32::new(0));
     let notify_on_send = Arc::new(Notify::new());
     let message_count = 100;