Add processor

Author: fpacifici

src/bin/consumer/asyncnoop.rs

@@ -1,21 +1,17 @@
 use clap::{App, Arg};
-use log::{debug, error, info};
+use log::{debug, info};
 use rdkafka::config::{ClientConfig, RDKafkaLogLevel};
 use rdkafka::consumer::stream_consumer::StreamConsumer;
-use rdkafka::consumer::CommitMode;
 use rdkafka::consumer::Consumer;
 use rdkafka::producer::FutureProducer;
 use rdkafka::util::get_rdkafka_version;
-use rust_arroyo::backends::kafka::create_kafka_message;
 use rust_arroyo::backends::AssignmentCallbacks;
-use rust_arroyo::processing::strategies::async_noop::build_topic_partitions;
 use rust_arroyo::processing::strategies::async_noop::AsyncNoopCommit;
 use rust_arroyo::processing::strategies::async_noop::CustomContext;
+use rust_arroyo::processing::StreamingStreamProcessor;
 use rust_arroyo::types::{Partition, Topic};
 use std::collections::HashMap;
-use std::time::Duration;
 use std::time::SystemTime;
-use tokio::time::timeout;
 
 // A type alias with your custom consumer can be created for convenience.
 type LoggingConsumer = StreamConsumer<CustomContext>;
@@ -71,7 +67,7 @@ async fn consume_and_produce(
     );
     let batch = Vec::new();
 
-    let mut strategy = AsyncNoopCommit {
+    let strategy = AsyncNoopCommit {
         topic: topic_clone,
         producer,
         batch,
@@ -80,35 +76,16 @@ async fn consume_and_produce(
         dest_topic: dest_topic.to_string(),
         source_topic: source_topic.to_string(),
     };
-    loop {
-        match timeout(Duration::from_secs(2), consumer.recv()).await {
-            Ok(result) => match result {
-                Err(e) => panic!("Kafka error: {}", e),
-                Ok(m) => {
-                    match strategy.poll().await {
-                        Some(partition_list) => {
-                            let part_list = build_topic_partitions(partition_list);
-                            consumer.commit(&part_list, CommitMode::Sync).unwrap();
-                            info!("Committed: {:?}", part_list);
-                        }
-                        None => {}
-                    }
 
-                    strategy.submit(create_kafka_message(m)).await;
-                }
-            },
-            Err(_) => {
-                error!("timeoout, flushing batch");
-                match strategy.poll().await {
-                    Some(partition_list) => {
-                        let part_list = build_topic_partitions(partition_list);
-                        consumer.commit(&part_list, CommitMode::Sync).unwrap();
-                        info!("Committed: {:?}", part_list);
-                    }
-                    None => {}
-                }
-            }
-        }
+    let mut processor = StreamingStreamProcessor {
+        consumer,
+        strategy,
+        shutdown_requested: false,
+    };
+
+    match processor.run().await {
+        Ok(_) => {}
+        Err(_) => panic!("Kafka error"),
     }
 }
 

src/processing/mod.rs

@@ -1,14 +1,23 @@
 pub mod strategies;
 
+use crate::backends::kafka::create_kafka_message;
 use crate::backends::{AssignmentCallbacks, Consumer};
+use crate::processing::strategies::async_noop::build_topic_partitions;
+use crate::processing::strategies::async_noop::AsyncNoopCommit;
+use crate::processing::strategies::async_noop::CustomContext;
 use crate::types::{Message, Partition, Topic};
 use async_mutex::Mutex;
 use futures::executor::block_on;
+use log::{error, info};
+use rdkafka::consumer::stream_consumer::StreamConsumer;
+use rdkafka::consumer::CommitMode;
+use rdkafka::consumer::Consumer as RdConsumer;
 use std::collections::HashMap;
 use std::mem::replace;
 use std::sync::Arc;
 use std::time::Duration;
 use strategies::{ProcessingStrategy, ProcessingStrategyFactory};
+use tokio::time::timeout;
 
 #[derive(Debug, Clone)]
 pub struct InvalidState;
@@ -210,6 +219,61 @@ impl<'a, TPayload: 'static + Clone> StreamProcessor<'a, TPayload> {
     }
 }
 
+pub struct StreamingStreamProcessor {
+    pub consumer: StreamConsumer<CustomContext>,
+    pub strategy: AsyncNoopCommit,
+    //message: Option<Message<TPayload>>,
+    pub shutdown_requested: bool,
+}
+
+impl StreamingStreamProcessor {
+    pub async fn run_once(&mut self) -> Result<(), RunError> {
+        match timeout(Duration::from_secs(2), self.consumer.recv()).await {
+            Ok(result) => match result {
+                Err(e) => panic!("Kafka error: {}", e),
+                Ok(m) => {
+                    match self.strategy.poll().await {
+                        Some(partition_list) => {
+                            let part_list = build_topic_partitions(partition_list);
+                            self.consumer.commit(&part_list, CommitMode::Sync).unwrap();
+                            info!("Committed: {:?}", part_list);
+                        }
+                        None => {}
+                    }
+
+                    self.strategy.submit(create_kafka_message(m)).await;
+                }
+            },
+            Err(_) => {
+                error!("timeoout, flushing batch");
+                match self.strategy.poll().await {
+                    Some(partition_list) => {
+                        let part_list = build_topic_partitions(partition_list);
+                        self.consumer.commit(&part_list, CommitMode::Sync).unwrap();
+                        info!("Committed: {:?}", part_list);
+                    }
+                    None => {}
+                }
+            }
+        }
+        Ok(())
+    }
+
+    /// The main run loop, see class docstring for more information.
+    pub async fn run(&mut self) -> Result<(), RunError> {
+        loop {
+            match self.run_once().await {
+                Ok(()) => {}
+                Err(e) => return Err(e),
+            }
+        }
+    }
+
+    pub fn signal_shutdown(&mut self) {
+        self.shutdown_requested = true;
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::strategies::{