Tweak Kafka source

src/channel.rs

@@ -138,7 +138,6 @@ impl<T: Send + 'static> Receiver<T> {
     }
 
     #[inline]
-
     pub async fn recv_async(&self) -> Result<T, RecvError> {
         self.0.recv_async().await.map_err(RecvError::from)
     }

src/operator/mod.rs

@@ -2161,7 +2161,6 @@ where
     ///
     /// **Warning**: [StreamCache] methods must only be called **after** the original `StreamContext`
     /// has finished executing. Calling `stream_in` or `inner_cloned` on an incomplete cache will panic!
-
     /// ## Example
     /// ```
     /// # use renoir::prelude::*;
@@ -2281,7 +2280,6 @@ where
     ///
     /// assert_eq!(res.get().unwrap(), vec![0, 1, 4, 9, 0, 1, 4, 9, 0, 1]);
     /// ```
-
     pub fn map_memo<O: Data + Sync, F>(
         self,
         f: F,

src/operator/source/kafka.rs

@@ -22,6 +22,7 @@ enum KafkaSourceInner {
     Running {
         rx: Receiver<OwnedMessage>,
         cancel_token: Arc<AtomicBool>,
+        cooldown: bool,
     }, // Terminated,
 }
 
@@ -81,6 +82,7 @@ impl Operator for KafkaSource {
         let (tx, rx) = flume::bounded(8);
         let cancel_token = Arc::new(AtomicBool::new(false));
         let cancel = cancel_token.clone();
+        tracing::debug!("started kafka source with topics {:?}", topics);
         tokio::spawn(async move {
             let mut stream = consumer.stream();
             while let Some(msg) = stream.next().await {
@@ -101,19 +103,43 @@ impl Operator for KafkaSource {
                     .expect("kafka fail to commit");
             }
         });
-        self.inner = KafkaSourceInner::Running { rx, cancel_token };
+        self.inner = KafkaSourceInner::Running {
+            rx,
+            cancel_token,
+            cooldown: false,
+        };
     }
 
     fn next(&mut self) -> StreamElement<Self::Out> {
-        match &self.inner {
+        match &mut self.inner {
             KafkaSourceInner::Init { .. } => {
                 unreachable!("KafkaSource executing before setup!")
             }
             // KafkaSourceInner::Terminated => return StreamElement::Terminate,
-            KafkaSourceInner::Running { rx, .. } => {
-                match rx.recv() {
+            KafkaSourceInner::Running { rx, cooldown, .. } => {
+                if *cooldown {
+                    match rx.recv() {
+                        Ok(msg) => {
+                            *cooldown = false;
+                            return StreamElement::Item(msg);
+                        }
+                        Err(flume::RecvError::Disconnected) => {
+                            tracing::warn!("kafka background task disconnected.");
+                            return StreamElement::Terminate;
+                        }
+                    }
+                }
+
+                match rx.recv_timeout(std::time::Duration::from_millis(100)) {
                     Ok(msg) => StreamElement::Item(msg),
-                    Err(e) => panic!("kafka background task panicked: {e}"),
+                    Err(flume::RecvTimeoutError::Timeout) => {
+                        *cooldown = true;
+                        StreamElement::FlushBatch
+                    }
+                    Err(flume::RecvTimeoutError::Disconnected) => {
+                        tracing::warn!("kafka background task disconnected.");
+                        StreamElement::Terminate
+                    }
                     // StreamElement::Terminate,
                 }
             }