Add recv to Arroyo Kafka consumer

Author: fpacifici

src/backends/kafka/mod.rs

@@ -209,6 +209,14 @@ impl<'a> ArroyoConsumer<'a, KafkaPayload> for KafkaConsumer {
         }
     }
 
+    async fn recv(&mut self) -> Result<ArroyoMessage<KafkaPayload>, ConsumerError> {
+        let consumer = self.consumer.as_mut().unwrap();
+        match consumer.recv().await {
+            Ok(result) => Ok(create_kafka_message(result)),
+            Err(e) => Err(ConsumerError::BrokerError(Box::new(e))),
+        }
+    }
+
     fn pause(&mut self, partitions: HashSet<Partition>) -> Result<(), ConsumerError> {
         self.state.assert_consuming_state()?;
 

src/backends/mod.rs

@@ -95,6 +95,8 @@ pub trait Consumer<'a, TPayload: Clone> {
         timeout: Option<Duration>,
     ) -> Result<Option<Message<TPayload>>, ConsumerError>;
 
+    async fn recv(&mut self) -> Result<Message<TPayload>, ConsumerError>;
+
     /// Pause consuming from the provided partitions.
     ///
     /// A partition that is paused will be automatically resumed during

src/bin/consumer/asyncnoop.rs

@@ -11,6 +11,7 @@ use rust_arroyo::types::{Partition, Topic};
 use std::collections::HashMap;
 use std::time::Duration;
 use std::time::SystemTime;
+use tokio::time::timeout;
 
 struct EmptyCallbacks {}
 impl AssignmentCallbacks for EmptyCallbacks {
@@ -64,30 +65,25 @@ async fn consume_and_produce(
         source_topic: source_topic.to_string(),
     };
     loop {
-        match consumer.poll(Some(Duration::from_secs(2))).await {
-            Ok(result) => match result {
-                None => {
-                    match strategy.poll().await {
-                        Some(request) => {
-                            consumer.stage_positions(request.positions).await.unwrap();
-                            consumer.commit_positions().await.unwrap();
-                            //info!("Committed: {:?}", request);
+        match timeout(Duration::from_secs(2), consumer.recv()).await {
+            Ok(result) => {
+                match result {
+                    Ok(message) => {
+                        match strategy.poll().await {
+                            Some(request) => {
+                                consumer.stage_positions(request.positions).await.unwrap();
+                                consumer.commit_positions().await.unwrap();
+                                //info!("Committed: {:?}", request);
+                            }
+                            None => {}
                         }
-                        None => {}
+                        strategy.submit(message).await;
                     }
-                }
-                Some(m) => {
-                    match strategy.poll().await {
-                        Some(request) => {
-                            consumer.stage_positions(request.positions).await.unwrap();
-                            consumer.commit_positions().await.unwrap();
-                            //info!("Committed: {:?}", request);
-                        }
-                        None => {}
+                    Err(e) => {
+                        panic!("Kafka error: {}", e)
                     }
-                    strategy.submit(m).await;
                 }
-            },
+            }
             Err(_) => {
                 error!("timeoout, flushing batch");
                 match strategy.poll().await {