super_stream_producer first basic version

Author: DanielePalaia

src/superstream.rs

@@ -9,20 +9,16 @@ pub enum RouteType {
     Key = 1,
 }
 
-trait Metadata {
-    async fn partitions(&mut self) -> Vec<String>;
-    async fn routes(&mut self, routing_key: String) -> Vec<String>;
+#[derive(Clone)]
+pub struct DefaultSuperStreamMetadata {
+    pub super_stream: String,
+    pub client: Client,
+    pub partitions: Vec<String>,
+    pub routes: Vec<String>,
 }
 
-struct DefaultSuperStreamMetadata {
-    super_stream: String,
-    client: Client,
-    partitions: Vec<String>,
-    routes: Vec<String>,
-}
-
-impl Metadata for DefaultSuperStreamMetadata {
-    async fn partitions(&mut self) -> Vec<String> {
+impl DefaultSuperStreamMetadata {
+    pub async fn partitions(&mut self) -> Vec<String> {
         if self.partitions.len() == 0 {
             let response = self.client.partitions(self.super_stream.clone()).await;
 
@@ -31,7 +27,7 @@ impl Metadata for DefaultSuperStreamMetadata {
 
         return self.partitions.clone();
     }
-    async fn routes(&mut self, routing_key: String) -> Vec<String> {
+    pub async fn routes(&mut self, routing_key: String) -> Vec<String> {
         if self.routes.len() == 0 {
             let response = self
                 .client
@@ -45,16 +41,17 @@ impl Metadata for DefaultSuperStreamMetadata {
     }
 }
 
-pub trait RoutingStrategy {
-    async fn routes(&self, message: Message, metadata: &mut impl Metadata) -> Vec<String>;
-}
-
-struct RoutingKeyRoutingStrategy {
+#[derive(Clone)]
+pub struct RoutingKeyRoutingStrategy {
     routing_extractor: &'static dyn Fn(Message) -> String,
 }
 
-impl RoutingStrategy for RoutingKeyRoutingStrategy {
-    async fn routes(&self, message: Message, metadata: &mut impl Metadata) -> Vec<String> {
+impl RoutingKeyRoutingStrategy {
+    pub async fn routes(
+        &self,
+        message: Message,
+        metadata: &mut DefaultSuperStreamMetadata,
+    ) -> Vec<String> {
         let key = (self.routing_extractor)(message);
 
         let routes = metadata.routes(key).await;
@@ -63,12 +60,17 @@ impl RoutingStrategy for RoutingKeyRoutingStrategy {
     }
 }
 
-struct HashRoutingMurmurStrategy {
+#[derive(Clone)]
+pub struct HashRoutingMurmurStrategy {
     routing_extractor: &'static dyn Fn(Message) -> String,
 }
 
-impl RoutingStrategy for HashRoutingMurmurStrategy {
-    async fn routes(&self, message: Message, metadata: &mut impl Metadata) -> Vec<String> {
+impl HashRoutingMurmurStrategy {
+    pub async fn routes(
+        &self,
+        message: Message,
+        metadata: &mut DefaultSuperStreamMetadata,
+    ) -> Vec<String> {
         let mut streams: Vec<String> = Vec::new();
 
         let key = (self.routing_extractor)(message);
@@ -83,3 +85,9 @@ impl RoutingStrategy for HashRoutingMurmurStrategy {
         return streams;
     }
 }
+
+#[derive(Clone)]
+pub enum RoutingStrategy {
+    HashRoutingStrategy(HashRoutingMurmurStrategy),
+    RoutingKeyStrategy(RoutingKeyRoutingStrategy),
+}

src/superstream_producer.rs

@@ -2,77 +2,85 @@ use crate::{
     client::Client,
     environment::Environment,
     error::{ProducerCreateError, ProducerPublishError},
-    producer::{ConfirmationStatus, Producer},
-    superstream::{RouteType, RoutingStrategy},
+    producer::{ConfirmationStatus, Producer, NoDedup},
+    superstream::{DefaultSuperStreamMetadata, RoutingStrategy},
 };
 use rabbitmq_stream_protocol::message::Message;
+use std::collections::HashMap;
 use std::future::Future;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
 type FilterValueExtractor = Arc<dyn Fn(&Message) -> String + 'static + Send + Sync>;
 
 #[derive(Clone)]
-pub struct NoDedup {}
-
-pub struct Dedup {}
-
-#[derive(Clone)]
-pub struct SuperStreamProducer<T>(Arc<SuperStreamProducerInternal>, PhantomData<T>);
+pub struct SuperStreamProducer<T>(Arc<SuperStreamProducerInternal<T>>, PhantomData<T>);
 
 /// Builder for [`SuperStreamProducer`]
-pub struct SuperStreamProducerBuilder<'a> {
+pub struct SuperStreamProducerBuilder<T> {
     pub(crate) environment: Environment,
     pub(crate) name: Option<String>,
     pub filter_value_extractor: Option<FilterValueExtractor>,
-    pub routing_strategy: &'a dyn RoutingStrategy,
+    pub routing_strategy: RoutingStrategy,
+    pub(crate) data: PhantomData<T>,
 }
 
-pub struct SuperStreamProducerInternal<'a> {
+pub struct SuperStreamProducerInternal<T> {
+    pub(crate) environment: Environment,
     client: Client,
     super_stream: String,
     publish_version: u16,
-    producers: Vec<Producer()>,
+    producers: HashMap<String, Producer<T>>,
     filter_value_extractor: Option<FilterValueExtractor>,
-    routing_strategy: &'a dyn RoutingStrategy,
+    super_stream_metadata: DefaultSuperStreamMetadata,
+    routing_strategy: RoutingStrategy,
 }
 
 impl SuperStreamProducer<NoDedup> {
     pub async fn send<Fut>(
         &self,
         message: Message,
-        cb: impl Fn(Result<ConfirmationStatus, ProducerPublishError>) -> Fut + Send + Sync + 'static,
-    ) -> Result<(), ProducerPublishError>
-    where
+        cb: impl Fn(Result<ConfirmationStatus, ProducerPublishError>) -> Fut + Send + Sync + 'static + Clone,
+    ) where
         Fut: Future<Output = ()> + Send + Sync + 'static,
     {
-        self.do_send(message, cb).await
-    }
-}
+        let routes = match self.0.routing_strategy.clone() {
+            RoutingStrategy::HashRoutingStrategy(routing_strategy) => {
+                routing_strategy
+                    .routes(message.clone(), &mut self.0.super_stream_metadata.clone())
+                    .await
+            }
+            RoutingStrategy::RoutingKeyStrategy(routing_strategy) => {
+                routing_strategy
+                    .routes(message.clone(), &mut self.0.super_stream_metadata.clone())
+                    .await
+            }
+        };
 
-impl SuperStreamProducer<Dedup> {
-    pub async fn send<Fut>(
-        &self,
-        message: Message,
-        cb: impl Fn(Result<ConfirmationStatus, ProducerPublishError>) -> Fut + Send + Sync + 'static,
-    ) -> Result<(), ProducerPublishError>
-    where
-        Fut: Future<Output = ()> + Send + Sync + 'static,
-    {
-        self.do_send(message, cb).await
+        for route in routes.into_iter() {
+            if !self.0.producers.contains_key(route.as_str()) {
+                let producer = self.0.environment.producer().build(route.as_str()).await;
+                self.0.producers.clone().insert(route.clone(), producer.unwrap().clone());
+            }
+
+            let producer = self.0.producers.get(route.as_str()).unwrap();
+            _ = producer
+                .send(message.clone(), cb.clone())
+                .await;
+        }
     }
 }
 
 impl<T> SuperStreamProducerBuilder<T> {
     pub async fn build(
         self,
         super_stream: &str,
-        route_type: RouteType,
+        route_type: RoutingStrategy,
     ) -> Result<SuperStreamProducer<T>, ProducerCreateError> {
         // Connect to the user specified node first, then look for the stream leader.
         // The leader is the recommended node for writing, because writing to a replica will redundantly pass these messages
         // to the leader anyway - it is the only one capable of writing.
-        let mut client = self.environment.create_client().await?;
+        let client = self.environment.create_client().await?;
 
         let mut publish_version = 1;
 
@@ -84,14 +92,21 @@ impl<T> SuperStreamProducerBuilder<T> {
             }
         }
 
-        let producers = Vec::new();
+        let producers = HashMap::new();
 
         let super_stream_producer = SuperStreamProducerInternal {
+            environment: self.environment.clone(),
             super_stream: super_stream.to_string(),
             client,
             publish_version,
             filter_value_extractor: self.filter_value_extractor,
             routing_strategy: route_type,
+            super_stream_metadata: DefaultSuperStreamMetadata {
+                super_stream: super_stream.to_string(),
+                client: self.environment.create_client().await?,
+                partitions: Vec::new(),
+                routes: Vec::new(),
+            },
             producers,
         };