implement basic kafka producer

Author: matthieugouel

src/consumer.rs

@@ -1,49 +1,19 @@
-use log::{info, warn};
+use log::warn;
 
-use rdkafka::client::ClientContext;
 use rdkafka::config::{ClientConfig, RDKafkaLogLevel};
 use rdkafka::consumer::stream_consumer::StreamConsumer;
-use rdkafka::consumer::{BaseConsumer, CommitMode, Consumer, ConsumerContext, Rebalance};
-use rdkafka::error::KafkaResult;
-use rdkafka::topic_partition_list::TopicPartitionList;
-
-use crate::handler;
-
-// A context can be used to change the behavior of producers and consumers by adding callbacks
-// that will be executed by librdkafka.
-// This particular context sets up custom callbacks to log rebalancing events.
-struct CustomContext;
-
-impl ClientContext for CustomContext {}
-
-impl ConsumerContext for CustomContext {
-    fn pre_rebalance(&self, _: &BaseConsumer<Self>, rebalance: &Rebalance) {
-        info!("Pre rebalance {:?}", rebalance);
-    }
-
-    fn post_rebalance(&self, _: &BaseConsumer<Self>, rebalance: &Rebalance) {
-        info!("Post rebalance {:?}", rebalance);
-    }
-
-    fn commit_callback(&self, result: KafkaResult<()>, _offsets: &TopicPartitionList) {
-        info!("Committing offsets: {:?}", result);
-    }
-}
-
-// A type alias with your custom consumer can be created for convenience.
-type LoggingConsumer = StreamConsumer<CustomContext>;
+use rdkafka::consumer::{CommitMode, Consumer, DefaultConsumerContext};
 
+#[allow(dead_code)]
 pub async fn consume(brokers: &str, group_id: &str, topics: &[&str]) {
-    let context = CustomContext;
+    let context = DefaultConsumerContext;
 
-    let consumer: LoggingConsumer = ClientConfig::new()
+    let consumer: StreamConsumer<DefaultConsumerContext> = ClientConfig::new()
         .set("group.id", group_id)
         .set("bootstrap.servers", brokers)
         .set("enable.partition.eof", "false")
         .set("session.timeout.ms", "6000")
         .set("enable.auto.commit", "true")
-        //.set("statistics.interval.ms", "30000")
-        //.set("auto.offset.reset", "smallest")
         .set_log_level(RDKafkaLogLevel::Debug)
         .create_with_context(context)
         .expect("Consumer creation failed");
@@ -56,8 +26,8 @@ pub async fn consume(brokers: &str, group_id: &str, topics: &[&str]) {
         match consumer.recv().await {
             Err(e) => warn!("Kafka error: {}", e),
             Ok(m) => {
-                let _ = handler::handle(&m).await;
-                consumer.commit_message(&m, CommitMode::Async).unwrap();
+                // Do something with the message
+                let _ = consumer.commit_message(&m, CommitMode::Async).unwrap();
             }
         };
     }

src/handler.rs

@@ -1,16 +1,17 @@
 use anyhow::Result;
 use caracat::models::Probe;
 use caracat::rate_limiter::RateLimitingMethod;
-use log::{info, warn};
-use rdkafka::message::BorrowedMessage;
-use rdkafka::message::{Headers, Message};
+// use log::{info, warn};
+// use rdkafka::message::BorrowedMessage;
+// use rdkafka::message::{Headers, Message};
 use std::net::IpAddr;
 use std::net::Ipv4Addr;
 use std::net::Ipv6Addr;
 use std::time::Duration;
 use tokio::task;
 
 use crate::prober::probe;
+use crate::producer::produce;
 
 /// Probing configuration.
 #[derive(Debug)]
@@ -79,31 +80,39 @@ fn decode_payload(payload: &str) -> Result<Payload> {
     })
 }
 
-pub async fn handle(m: &BorrowedMessage<'_>) -> Result<()> {
-    let payload = match m.payload_view::<str>() {
-        None => "",
-        Some(Ok(s)) => s,
-        Some(Err(e)) => {
-            warn!("Error while deserializing message payload: {:?}", e);
-            ""
-        }
-    };
+// pub async fn handle(m: &BorrowedMessage<'_>) -> Result<()> {
+pub async fn handle(
+    brokers: &str,
+    _in_group_id: &str,
+    _in_topics: &[&str],
+    out_topic: &str,
+) -> Result<()> {
+    // let payload = match m.payload_view::<str>() {
+    //     None => "",
+    //     Some(Ok(s)) => s,
+    //     Some(Err(e)) => {
+    //         warn!("Error while deserializing message payload: {:?}", e);
+    //         ""
+    //     }
+    // };
 
-    info!(
-        "key: '{:?}', payload: '{}', topic: {}, partition: {}, offset: {}, timestamp: {:?}",
-        m.key(),
-        payload,
-        m.topic(),
-        m.partition(),
-        m.offset(),
-        m.timestamp()
-    );
+    // info!(
+    //     "key: '{:?}', payload: '{}', topic: {}, partition: {}, offset: {}, timestamp: {:?}",
+    //     m.key(),
+    //     payload,
+    //     m.topic(),
+    //     m.partition(),
+    //     m.offset(),
+    //     m.timestamp()
+    // );
 
-    if let Some(headers) = m.headers() {
-        for header in headers.iter() {
-            info!("  Header {:#?}: {:?}", header.key, header.value);
-        }
-    }
+    // if let Some(headers) = m.headers() {
+    //     for header in headers.iter() {
+    //         info!("  Header {:#?}: {:?}", header.key, header.value);
+    //     }
+    // }
+
+    let payload = "2606:4700:4700::1111,1,32";
 
     // Probing
     let config = create_config();
@@ -119,18 +128,10 @@ pub async fn handle(m: &BorrowedMessage<'_>) -> Result<()> {
             protocol: caracat::models::L4::ICMPv6,
         });
     }
-    let result = task::spawn_blocking(move || {
-        probe(
-            config,
-            probes_to_send.into_iter(),
-            Some(String::from("./test.csv")),
-        )
-    })
-    .await?;
+    let result = task::spawn_blocking(move || probe(config, probes_to_send.into_iter())).await?;
 
-    if let Err(e) = result {
-        warn!("Error while probing: {:?}", e);
-    }
+    let (_, _, results) = result?;
+    produce(brokers, out_topic, results).await;
 
     Ok(())
 }

src/main.rs

@@ -1,6 +1,7 @@
 mod consumer;
 mod handler;
 mod prober;
+mod producer;
 
 use anyhow::Result;
 use chrono::Local;
@@ -9,7 +10,8 @@ use clap_verbosity_flag::{InfoLevel, Verbosity};
 use env_logger::Builder;
 use std::io::Write;
 
-use crate::consumer::consume;
+// use crate::consumer::consume;
+use crate::handler::handle;
 
 #[derive(CliParser, Debug, Clone)]
 #[command(version, about, long_about = None)]
@@ -40,10 +42,12 @@ async fn main() -> Result<()> {
     set_logging(&cli);
 
     let brokers = "localhost:9092";
-    let group_id = "osiris";
-    let topics = vec!["osiris"];
+    let in_group_id = "osiris";
+    let in_topics = vec!["osiris"];
+    let out_topic = "results";
 
-    consume(brokers, group_id, &topics).await;
+    // consume(brokers, group_id, &topics).await;
+    let _ = handle(&brokers, &in_group_id, &in_topics, &out_topic).await;
 
     Ok(())
 }

src/prober.rs

@@ -4,7 +4,6 @@ use hyperloglog::HyperLogLog;
 
 use std::fmt::Display;
 use std::fmt::Formatter;
-use std::io::{stdout, Write};
 use std::sync::{Arc, Mutex};
 use std::thread;
 use std::thread::sleep;
@@ -13,7 +12,7 @@ use std::thread::JoinHandle;
 use log::{error, info, trace};
 
 use crate::handler::Config;
-use caracat::models::Probe;
+use caracat::models::{Probe, Reply};
 use caracat::rate_limiter::RateLimiter;
 use caracat::receiver::Receiver;
 use caracat::sender::Sender;
@@ -22,8 +21,11 @@ use caracat::sender::Sender;
 pub fn probe<T: Iterator<Item = Probe>>(
     config: Config,
     probes: T,
-    csv_output: Option<String>,
-) -> Result<(SendStatistics, ReceiveStatistics)> {
+) -> Result<(
+    SendStatistics,
+    ReceiveStatistics,
+    Arc<Mutex<Vec<caracat::models::Reply>>>,
+)> {
     info!("{:?}", config);
 
     let rate_limiter = RateLimiter::new(
@@ -36,7 +38,6 @@ pub fn probe<T: Iterator<Item = Probe>>(
         config.interface.clone(),
         config.instance_id,
         config.integrity_check,
-        csv_output,
     );
 
     let mut prober = SendLoop::new(
@@ -49,8 +50,8 @@ pub fn probe<T: Iterator<Item = Probe>>(
         rate_limiter,
         Sender::new(
             &config.interface,
-            None,
-            None,
+            config.src_ipv4_addr,
+            config.src_ipv6_addr,
             config.instance_id,
             config.dry_run,
         )?,
@@ -63,27 +64,26 @@ pub fn probe<T: Iterator<Item = Probe>>(
     // TODO: Cleaner way?
     let final_prober_statistics = *prober.statistics().lock().unwrap();
     let final_receiver_statistics = receiver.statistics().lock().unwrap().clone();
+    let results = Arc::clone(&receiver.results);
 
     receiver.stop();
 
-    Ok((final_prober_statistics, final_receiver_statistics))
+    Ok((
+        final_prober_statistics,
+        final_receiver_statistics,
+        Arc::clone(&results),
+    ))
 }
 
-// The pcap crate doesn't support `pcap_loop` and `pcap_breakloop`,
-// so we implement our own looping mechanism.
 pub struct ReceiveLoop {
     handle: JoinHandle<()>,
     stopped: Arc<Mutex<bool>>,
     statistics: Arc<Mutex<ReceiveStatistics>>,
+    results: Arc<Mutex<Vec<Reply>>>,
 }
 
 impl ReceiveLoop {
-    pub fn new(
-        interface: String,
-        instance_id: u16,
-        integrity_check: bool,
-        output_csv: Option<String>,
-    ) -> Self {
+    pub fn new(interface: String, instance_id: u16, integrity_check: bool) -> Self {
         // By default if a thread panic, the other threads are not affected and the error
         // is only surfaced when joining the thread. However since this is a long-lived thread,
         // we're not calling join until the end of the process. Since this loop is critical to
@@ -94,20 +94,23 @@ impl ReceiveLoop {
         let statistics = Arc::new(Mutex::new(ReceiveStatistics::default()));
         let statistics_thr = statistics.clone();
 
+        let results = Arc::new(Mutex::new(vec![]));
+        let results_thr = results.clone();
+
         let handle = thread::spawn(move || {
             let mut receiver = Receiver::new_batch(&interface).unwrap();
 
-            let wtr: Box<dyn Write> = match output_csv {
-                Some(output_csv) => {
-                    let file = std::fs::File::create(output_csv).unwrap();
-                    Box::new(std::io::BufWriter::new(file))
-                }
-                None => Box::new(stdout().lock()),
-            };
+            // let wtr: Box<dyn Write> = match output_csv {
+            //     Some(output_csv) => {
+            //         let file = std::fs::File::create(output_csv).unwrap();
+            //         Box::new(std::io::BufWriter::new(file))
+            //     }
+            //     None => Box::new(stdout().lock()),
+            // };
 
-            let mut csv_writer = csv::WriterBuilder::new()
-                .has_headers(false) // TODO: Set to true, but how to serialize MPLS labels?
-                .from_writer(wtr);
+            // let mut csv_writer = csv::WriterBuilder::new()
+            //     .has_headers(false) // TODO: Set to true, but how to serialize MPLS labels?
+            //     .from_writer(wtr);
 
             loop {
                 // TODO: Cleanup this loop & statistics handling
@@ -117,6 +120,9 @@ impl ReceiveLoop {
                 statistics.pcap_received = pcap_statistics.received;
                 statistics.pcap_dropped = pcap_statistics.dropped;
                 statistics.pcap_if_dropped = pcap_statistics.if_dropped;
+
+                let mut results = results_thr.lock().unwrap();
+
                 match result {
                     Ok(reply) => {
                         trace!("{:?}", reply);
@@ -130,7 +136,8 @@ impl ReceiveLoop {
                                     .icmp_messages_excl_dest
                                     .insert(&reply.reply_src_addr);
                             }
-                            csv_writer.serialize(reply).unwrap();
+
+                            results.push(reply);
                             // TODO: Write round column.
                             // TODO: Compare output with caracal (capture timestamp resolution?)
                         } else {
@@ -158,12 +165,13 @@ impl ReceiveLoop {
                     break;
                 }
             }
-            csv_writer.flush().unwrap();
         });
+
         ReceiveLoop {
             handle,
             stopped,
             statistics,
+            results,
         }
     }
 
@@ -271,7 +279,6 @@ impl SendLoop {
     }
 }
 
-
 #[derive(Clone, Debug)]
 pub struct ReceiveStatistics {
     /// Number of packets received.