Add count

Cargo.toml

@@ -15,6 +15,8 @@ rand = {  version = "0.10.0", features = ["thread_rng"]}
 rand_distr = "0.6.0"
 flatbuffers = "25.12.19"
 serde_json = "1.0.149"
+tokio = { version = "1.51.0", features = ["full"] }
+futures = "0.3.32"
 
 [target.'cfg(windows)'.dependencies]
 rdkafka = { version="0.39.0", features = ["cmake_build"] }

README.md

@@ -63,6 +63,15 @@ INFO:saluki:            0 - low:7515, high:7551, num_messages:36
 
 `saluki play mybroker:9092/source_topic mybroker:9092/dest_topic -t 1762209990 1762209992` - This will forward messages between the two given timestamps.
 
+
+## `count` - Count topic data rates
+
+`count` is used for viewing the current data rate in a given topic. 
+
+An example of using this could be:
+
+`saluki count mybroker:9092/mytopic --message-interval 3` - this listens to the `mytopic` topic and prints the data rate between 3 second intervals. 
+
 # Developer setup
 
 some system dependencies are required. On Windows these are built-in, but on a debian-based linux distro you will need `libcurl4-openssl-dev`

src/consume.rs

@@ -79,7 +79,7 @@ pub fn consume(
     }
 
     consumer
-        .subscribe(&[&*topic.topic])
+        .subscribe(&[&topic.topic])
         .unwrap_or_else(|_| panic!("Failed to subscribe to topic {}", topic.topic));
 
     let mut counter = 0;
@@ -119,7 +119,7 @@ pub fn consume(
         }
         counter += 1;
         if Some(counter) == num_messages || Some(counter) == last {
-            info!("Reached {} messages, exiting", counter);
+            println!("Reached {} messages, exiting", counter);
             break;
         }
     }

src/count.rs

@@ -0,0 +1,51 @@
+use crate::cli_utils::BrokerAndTopic;
+use futures::stream::StreamExt;
+use log::error;
+use rdkafka::consumer::{Consumer, DefaultConsumerContext, StreamConsumer};
+use rdkafka::{ClientConfig, Message};
+use tokio::time::{self, Duration};
+use uuid::Uuid;
+
+pub async fn count(topic: BrokerAndTopic, message_interval: u64) {
+    let consumer: StreamConsumer<DefaultConsumerContext> = ClientConfig::new()
+        .set("group.id", Uuid::new_v4().to_string())
+        .set("bootstrap.servers", topic.broker())
+        .create()
+        .expect("Consumer creation failed");
+
+    consumer
+        .subscribe(&[&topic.topic])
+        .unwrap_or_else(|e| panic!("Failed to subscribe to topic {}: {}", topic.topic, e));
+
+    let start = std::time::Instant::now();
+    let mut bytes_this_second: usize = 0;
+    let mut total_bytes: usize = 0;
+    let mut stream = consumer.stream();
+    let mut interval = time::interval(Duration::from_secs(message_interval));
+
+    loop {
+        tokio::select! {
+            _msg = stream.next() => {
+                match _msg {
+                    Some(Ok(msg)) => {
+                        if msg.payload().is_some() {
+                            bytes_this_second += msg.payload_len();
+                            total_bytes += msg.payload_len();
+                        }
+
+                    },
+                    Some(Err(e)) => error!("Error reading from stream {:?}", e),
+                    None => {}
+                }
+            }
+            _ = interval.tick() => {
+                println!("{:.5} Mbit/s (since program start: average {:.5} Mbit/s, {:.5} MB total)",
+                bytes_this_second as f64/125000.0,
+                total_bytes as f64 / 125000.0 / start.elapsed().as_secs_f64(),
+                total_bytes as f64 / 1_000_000.0
+            );
+                bytes_this_second = 0;
+            }
+        }
+    }
+}