Dynamic send (#276)

* Implement dynamic send

Author: allevo

.tarpaulin.toml

@@ -4,6 +4,7 @@ exclude-files=[
     "benchmark/*",
     "build.rs",
     "src/lib.rs",
+    "src/bin/**/*.rs",
     "tests/**/*",
     "mod.rs"
 ]

Cargo.toml

@@ -45,6 +45,8 @@ tracing-subscriber = "0.3.1"
 fake = { version = "3.0.0", features = ['derive'] }
 chrono = "0.4.26"
 serde_json = "1.0"
+reqwest = { version = "0.12", features = ["json"] }
+serde = { version = "1.0", features = ["derive"] }
 
 [features]
 default = []
@@ -66,3 +68,6 @@ path="examples/superstreams/send_super_stream.rs"
 name="environment_deserialization"
 path="examples/environment_deserialization.rs"
 
+[[bin]]
+name = "perf-producer"
+path = "src/bin/perf-producer.rs"

examples/send_async.rs

@@ -32,7 +32,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let create_response = environment
         .stream_creator()
         .max_length(ByteCapacity::GB(5))
-        .create(&stream)
+        .create(stream)
         .await;
 
     if let Err(e) = create_response {

examples/tls_producer.rs

@@ -54,9 +54,9 @@ async fn start_publisher(
     env: Environment,
     stream: &String,
 ) -> Result<(), Box<dyn std::error::Error>> {
-    let _ = env.stream_creator().create(&stream).await;
+    let _ = env.stream_creator().create(stream).await;
 
-    let producer = env.producer().batch_size(BATCH_SIZE).build(&stream).await?;
+    let producer = env.producer().batch_size(BATCH_SIZE).build(stream).await?;
 
     let is_batch_send = true;
     tokio::task::spawn(async move {

protocol/src/message/builder.rs

@@ -30,7 +30,7 @@ impl MessageBuilder {
     pub fn application_properties(self) -> ApplicationPropertiesBuider {
         ApplicationPropertiesBuider(self)
     }
-    pub fn publising_id(mut self, publishing_id: u64) -> Self {
+    pub fn publishing_id(mut self, publishing_id: u64) -> Self {
         self.0.publishing_id = Some(publishing_id);
         self
     }

src/bin/perf-producer.rs

@@ -0,0 +1,175 @@
+#![allow(dead_code)]
+
+use std::{
+    sync::{
+        atomic::{AtomicU32, Ordering},
+        Arc,
+    },
+    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
+};
+
+use rabbitmq_stream_client::{
+    types::{ByteCapacity, Message, OffsetSpecification},
+    Environment,
+};
+use tokio::{sync::mpsc::UnboundedSender, time::sleep};
+use tokio_stream::StreamExt;
+
+static ONE_SECOND: Duration = Duration::from_secs(1);
+static ONE_MINUTE: Duration = Duration::from_secs(60);
+
+struct Metric {
+    created_at: u128,
+    received_at: SystemTime,
+}
+
+#[derive(Debug)]
+struct Stats {
+    average_latency: f32,
+    messages_received: usize,
+}
+
+#[tokio::main]
+async fn main() {
+    let stream_name = "perf-stream";
+
+    let environment = Environment::builder().build().await.unwrap();
+    let _ = environment.delete_stream(stream_name).await;
+    environment
+        .stream_creator()
+        .max_length(ByteCapacity::GB(5))
+        .create(stream_name)
+        .await
+        .unwrap();
+
+    let environment = Arc::new(environment);
+
+    let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel();
+    let consumer_env = environment.clone();
+    let consumer_handler = tokio::spawn(async move {
+        start_consumer(consumer_env, stream_name, sender).await;
+    });
+
+    let produced_messages = AtomicU32::new(0);
+    let producer_env = environment.clone();
+    let producer_handler = tokio::spawn(async move {
+        start_producer(producer_env, stream_name, &produced_messages).await;
+    });
+
+    let run_for = Duration::from_secs(5 * 60);
+
+    tokio::spawn(async move {
+        sleep(run_for).await;
+        producer_handler.abort();
+        sleep(Duration::from_secs(1)).await;
+        consumer_handler.abort();
+    });
+
+    let minutes = run_for.as_secs() / 60;
+
+    let mut now = Instant::now();
+    // 5 minutes of metrics
+    let mut metrics = Vec::with_capacity(50 * 60 * minutes as usize);
+    while let Some(metric) = receiver.recv().await {
+        if now.elapsed() > ONE_MINUTE {
+            now = Instant::now();
+
+            let last_metrics = metrics;
+            metrics = Vec::with_capacity(50 * 60 * minutes as usize);
+            tokio::spawn(async move {
+                let stats = calculate_stats(last_metrics).await;
+                println!("stats: {:?}", stats);
+            });
+        }
+        metrics.push(metric);
+    }
+
+    let stats = calculate_stats(metrics).await;
+    println!("stats: {:?}", stats);
+}
+
+async fn calculate_stats(metrics: Vec<Metric>) -> Stats {
+    let mut total_latency = 0;
+    let metric_count = metrics.len();
+    for metric in metrics {
+        let created_at = SystemTime::UNIX_EPOCH + Duration::from_millis(metric.created_at as u64);
+        let received_at = metric.received_at;
+        let delta = received_at.duration_since(created_at).unwrap();
+        total_latency += delta.as_millis();
+    }
+
+    Stats {
+        average_latency: total_latency as f32 / metric_count as f32,
+        messages_received: metric_count,
+    }
+}
+
+async fn start_consumer(
+    environment: Arc<Environment>,
+    stream_name: &str,
+    sender: UnboundedSender<Metric>,
+) {
+    let mut consumer = environment
+        .consumer()
+        .offset(OffsetSpecification::First)
+        .build(stream_name)
+        .await
+        .unwrap();
+    while let Some(Ok(delivery)) = consumer.next().await {
+        let produced_at = delivery
+            .message()
+            .data()
+            .map(|data| {
+                u128::from_be_bytes([
+                    data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7],
+                    data[8], data[9], data[10], data[11], data[12], data[13], data[14], data[15],
+                ])
+            })
+            .unwrap();
+        let metric = Metric {
+            created_at: produced_at,
+            received_at: SystemTime::now(),
+        };
+        sender.send(metric).unwrap();
+    }
+}
+
+async fn start_producer(
+    environment: Arc<Environment>,
+    stream_name: &str,
+    produced_messages: &AtomicU32,
+) {
+    let message_per_second = 50_usize;
+    let producer = environment.producer().build(stream_name).await.unwrap();
+
+    loop {
+        let start = Instant::now();
+        let messages = create_messages(message_per_second);
+        let messages_sent = messages.len() as u32;
+        for message in messages {
+            producer.send(message, |_| async {}).await.unwrap();
+        }
+        produced_messages.fetch_add(messages_sent, Ordering::Relaxed);
+
+        let elapsed = start.elapsed();
+
+        if ONE_SECOND > elapsed {
+            sleep(ONE_SECOND - elapsed).await;
+        }
+    }
+}
+
+fn create_messages(message_count_per_batch: usize) -> Vec<Message> {
+    (0..message_count_per_batch)
+        .map(|_| {
+            let start = SystemTime::now();
+            let since_the_epoch = start
+                .duration_since(UNIX_EPOCH)
+                .expect("Time went backwards");
+            let since_the_epoch = since_the_epoch.as_millis();
+            Message::builder()
+                .body(since_the_epoch.to_be_bytes())
+                .build()
+        })
+        .collect()
+}

src/client/mod.rs

@@ -25,7 +25,7 @@ use tokio::{net::TcpStream, sync::Notify};
 use tokio_rustls::client::TlsStream;
 
 use tokio_util::codec::Framed;
-use tracing::trace;
+use tracing::{trace, warn};
 
 use crate::{error::ClientError, RabbitMQStreamResult};
 pub use message::ClientMessage;
@@ -275,8 +275,12 @@ impl Client {
         state.handler = Some(Arc::new(handler));
     }
 
+    pub fn is_closed(&self) -> bool {
+        self.channel.is_closed()
+    }
+
     pub async fn close(&self) -> RabbitMQStreamResult<()> {
-        if self.channel.is_closed() {
+        if self.is_closed() {
             return Err(ClientError::AlreadyClosed);
         }
         let _: CloseResponse = self
@@ -286,12 +290,17 @@ impl Client {
             .await?;
 
         let mut state = self.state.write().await;
-
+        // This stop the tokio task that performs heartbeats
         state.heartbeat_task.take();
-
         drop(state);
+
+        self.force_drop_connection().await
+    }
+
+    async fn force_drop_connection(&self) -> RabbitMQStreamResult<()> {
         self.channel.close().await
     }
+
     pub async fn subscribe(
         &self,
         subscription_id: u8,
@@ -622,7 +631,7 @@ impl Client {
         M: FnOnce(u32) -> R,
     {
         let Some((correlation_id, mut receiver)) = self.dispatcher.response_channel() else {
-            trace!("Connection si closed here");
+            trace!("Connection is closed here");
             return Err(ClientError::ConnectionClosed);
         };
 
@@ -711,9 +720,16 @@ impl Client {
         let heartbeat_task = tokio::spawn(async move {
             loop {
                 trace!("Sending heartbeat");
-                let _ = channel.send(HeartBeatCommand::default().into()).await;
+                if channel
+                    .send(HeartBeatCommand::default().into())
+                    .await
+                    .is_err()
+                {
+                    break;
+                }
                 tokio::time::sleep(Duration::from_secs(heartbeat_interval.into())).await;
             }
+            warn!("Heartbeat task stopped. Force closing connection");
         })
         .into();
         state.heartbeat_task = Some(heartbeat_task);

src/client/options.rs

@@ -223,6 +223,11 @@ impl ClientOptionsBuilder {
         self
     }
 
+    pub fn client_provided_name(mut self, client_provided_name: String) -> Self {
+        self.0.client_provided_name = client_provided_name;
+        self
+    }
+
     pub fn collector(mut self, collector: Arc<dyn MetricsCollector>) -> Self {
         self.0.collector = collector;
         self
@@ -587,7 +592,7 @@ mod tests {
         assert_eq!(options.heartbeat, 10000);
         assert_eq!(options.max_frame_size, 1);
         assert!(matches!(options.tls, TlsConfiguration::Untrusted));
-        assert_eq!(options.load_balancer_mode, true);
+        assert!(options.load_balancer_mode);
     }
 
     #[cfg(feature = "serde")]

src/environment.rs

@@ -1,6 +1,5 @@
 use std::marker::PhantomData;
 use std::sync::Arc;
-use std::time::Duration;
 
 use crate::types::OffsetSpecification;
 use crate::{client::TlsConfiguration, producer::NoDedup};
@@ -197,7 +196,6 @@ impl Environment {
             environment: self.clone(),
             name: None,
             batch_size: 100,
-            batch_publishing_delay: Duration::from_millis(100),
             data: PhantomData,
             filter_value_extractor: None,
             client_provided_name: String::from("rust-stream-producer"),
@@ -328,6 +326,7 @@ impl EnvironmentBuilder {
         self.0.client_options.heartbeat = heartbeat;
         self
     }
+
     pub fn metrics_collector(
         mut self,
         collector: impl MetricsCollector + 'static,

src/error.rs

@@ -109,6 +109,8 @@ pub enum ProducerPublishError {
     Confirmation { stream: String },
     #[error(transparent)]
     Client(#[from] ClientError),
+    #[error("Failed to publish message, timeout")]
+    Timeout,
 }
 
 #[derive(Error, Debug)]