implement filtering on super_stream

Author: DanielePalaia

src/environment.rs

@@ -59,7 +59,7 @@ impl Environment {
         SuperStreamProducerBuilder {
             environment: self.clone(),
             data: PhantomData,
-            //filter_value_extractor: None,
+            filter_value_extractor: None,
             route_strategy: routing_strategy,
         }
     }
@@ -78,6 +78,7 @@ impl Environment {
         SuperStreamConsumerBuilder {
             environment: self.clone(),
             offset_specification: OffsetSpecification::Next,
+            filter_configuration: None,
         }
     }
 

src/producer.rs

@@ -242,6 +242,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 {

src/superstream_consumer.rs

@@ -1,7 +1,7 @@
 use crate::consumer::Delivery;
 use crate::error::{ConsumerCloseError, ConsumerDeliveryError};
 use crate::superstream::DefaultSuperStreamMetadata;
-use crate::{error::ConsumerCreateError, ConsumerHandle, Environment};
+use crate::{error::ConsumerCreateError, ConsumerHandle, Environment, FilterConfiguration};
 use futures::task::AtomicWaker;
 use futures::{Stream, StreamExt};
 use rabbitmq_stream_protocol::commands::subscribe::OffsetSpecification;
@@ -13,8 +13,6 @@ use std::task::{Context, Poll};
 use tokio::sync::mpsc::{channel, Receiver};
 use tokio::task;
 
-//type FilterPredicate = Option<Arc<dyn Fn(&Message) -> bool + Send + Sync>>;
-
 /// API for consuming RabbitMQ stream messages
 pub struct SuperStreamConsumer {
     internal: Arc<SuperStreamConsumerInternal>,
@@ -25,12 +23,14 @@ struct SuperStreamConsumerInternal {
     closed: Arc<AtomicBool>,
     handlers: Vec<ConsumerHandle>,
     waker: AtomicWaker,
+    //filter_configuration: Option<FilterConfiguration>,
 }
 
 /// Builder for [`Consumer`]
 pub struct SuperStreamConsumerBuilder {
     pub(crate) environment: Environment,
     pub(crate) offset_specification: OffsetSpecification,
+    pub(crate) filter_configuration: Option<FilterConfiguration>,
 }
 
 impl SuperStreamConsumerBuilder {
@@ -58,6 +58,7 @@ impl SuperStreamConsumerBuilder {
                 .environment
                 .consumer()
                 .offset(self.offset_specification.clone())
+                .filter_input(self.filter_configuration.clone())
                 .build(partition.as_str())
                 .await
                 .unwrap();
@@ -75,6 +76,7 @@ impl SuperStreamConsumerBuilder {
             closed: Arc::new(AtomicBool::new(false)),
             handlers,
             waker: AtomicWaker::new(),
+            //filter_configuration: self.filter_configuration.clone(),
         };
 
         Ok(SuperStreamConsumer {
@@ -87,6 +89,11 @@ impl SuperStreamConsumerBuilder {
         self.offset_specification = offset_specification;
         self
     }
+
+    pub fn filter_input(mut self, filter_configuration: Option<FilterConfiguration>) -> Self {
+        self.filter_configuration = filter_configuration;
+        self
+    }
 }
 
 impl Stream for SuperStreamConsumer {

src/superstream_producer.rs

@@ -12,20 +12,21 @@ 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>(
     Arc<SuperStreamProducerInternal>,
     HashMap<String, Producer<T>>,
     DefaultSuperStreamMetadata,
     PhantomData<T>,
+    RoutingStrategy,
 );
 
 /// Builder for [`SuperStreamProducer`]
 pub struct SuperStreamProducerBuilder<T> {
     pub(crate) environment: Environment,
-    //pub filter_value_extractor: Option<FilterValueExtractor>,
+    pub filter_value_extractor: Option<FilterValueExtractor>,
     pub route_strategy: RoutingStrategy,
     pub(crate) data: PhantomData<T>,
 }
@@ -34,8 +35,8 @@ pub struct SuperStreamProducerInternal {
     pub(crate) environment: Environment,
     client: Client,
     // TODO: implement filtering for superstream
-    //filter_value_extractor: Option<FilterValueExtractor>,
-    routing_strategy: RoutingStrategy,
+    filter_value_extractor: Option<FilterValueExtractor>,
+    //routing_strategy: RoutingStrategy,
 }
 
 impl SuperStreamProducer<NoDedup> {
@@ -51,7 +52,7 @@ impl SuperStreamProducer<NoDedup> {
     where
         Fut: Future<Output = ()> + Send + Sync + 'static,
     {
-        let routes = match self.0.routing_strategy.clone() {
+        let routes = match self.4.clone() {
             RoutingStrategy::HashRoutingStrategy(routing_strategy) => {
                 routing_strategy.routes(&message, &mut self.2).await
             }
@@ -62,7 +63,13 @@ 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;
+                let producer = self
+                    .0
+                    .environment
+                    .producer()
+                    .filter_value_extractor_arc(self.0.filter_value_extractor.clone())
+                    .build(route.as_str())
+                    .await;
                 self.1.insert(route.clone(), producer.unwrap());
             }
 
@@ -115,8 +122,7 @@ impl<T> SuperStreamProducerBuilder<T> {
         let super_stream_producer = SuperStreamProducerInternal {
             environment: self.environment.clone(),
             client,
-            //filter_value_extractor: self.filter_value_extractor,
-            routing_strategy: self.route_strategy,
+            filter_value_extractor: self.filter_value_extractor,
         };
 
         let internal_producer = Arc::new(super_stream_producer);
@@ -125,8 +131,18 @@ impl<T> SuperStreamProducerBuilder<T> {
             producers,
             super_stream_metadata,
             PhantomData,
+            self.route_strategy,
         );
 
         Ok(super_stream_producer)
     }
+
+    pub fn filter_value_extractor(
+        mut self,
+        filter_value_extractor: impl Fn(&Message) -> String + Send + Sync + 'static,
+    ) -> Self {
+        let f = Arc::new(filter_value_extractor);
+        self.filter_value_extractor = Some(f);
+        self
+    }
 }

tests/integration/consumer_test.rs

@@ -14,7 +14,9 @@ use rabbitmq_stream_client::{
 use tokio::task;
 use tokio::time::sleep;
 
-use rabbitmq_stream_client::types::{HashRoutingMurmurStrategy, RoutingStrategy};
+use rabbitmq_stream_client::types::{
+    HashRoutingMurmurStrategy, RoutingKeyRoutingStrategy, RoutingStrategy,
+};
 use rabbitmq_stream_protocol::ResponseCode;
 use std::sync::atomic::{AtomicU32, Ordering};
 use {std::sync::Arc, std::sync::Mutex};

tests/integration/producer_test.rs

@@ -391,7 +391,7 @@ async fn producer_send_after_close_error() {
     );
 }
 
-fn routing_key_strategy_value_extractor(message: &Message) -> String {
+pub fn routing_key_strategy_value_extractor(message: &Message) -> String {
     return "0".to_string();
 }
 
@@ -459,7 +459,6 @@ async fn hash_super_steam_producer_test() {
         .unwrap();
 
     for i in 0..message_count {
-        println!("sending message {}", i);
         let counter = confirmed_messages.clone();
         let notifier = notify_on_send.clone();
         let msg = Message::builder().body(format!("message{}", i)).build();