feat(stream): add pull-based pipeline runtime [STREAM-1668]

Cargo.toml

@@ -16,6 +16,7 @@ path = "rust-arroyo/src/lib.rs"
 ssl = ["rdkafka/ssl"]
 
 [dependencies]
+async-stream = "0.3"
 chrono = "0.4.26"
 coarsetime = "0.1.33"
 once_cell = "1.18.0"
@@ -26,12 +27,15 @@ serde = { version = "1.0.137", features = ["derive"] }
 serde_json = "1.0.81"
 thiserror = "1.0"
 tokio = { version = "1.19.2", features = ["full"] }
+futures = "0.3"
+tokio-stream = { version = "0.1", features = ["time"] }
+tokio-util = "0.7"
 tracing = "0.1.40"
 uuid = { version = "1.5.0", features = ["v4"] }
 parking_lot = "0.12.1"
 
 [dev-dependencies]
-tracing-subscriber = "0.3.18"
+tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 
 [[example]]
 name = "base_processor"
@@ -44,3 +48,7 @@ path = "rust-arroyo/examples/transform_and_produce.rs"
 [[example]]
 name = "base_consumer"
 path = "rust-arroyo/examples/base_consumer.rs"
+
+[[example]]
+name = "transform_and_produce_pull"
+path = "rust-arroyo/examples/transform_and_produce_pull.rs"

rust-arroyo/examples/transform_and_produce_pull.rs

@@ -0,0 +1,97 @@
+/// Pull-based version of transform_and_produce.
+///
+/// Pipeline:
+///   KafkaSource → apply(reverse) → on_next(produce) → on_reject(log) → commit
+///
+/// `PipelineRunner::run()` handles the rebalance restart loop —
+/// the closure is called once per partition assignment with fresh
+/// stages, handlers, and tracker.
+extern crate sentry_arroyo;
+
+use std::time::Duration;
+
+use sentry_arroyo::backends::kafka::config::KafkaConfig;
+use sentry_arroyo::backends::kafka::producer::KafkaProducer;
+use sentry_arroyo::backends::kafka::types::KafkaPayload;
+use sentry_arroyo::backends::kafka::InitialOffset;
+use sentry_arroyo::processing::stream::{
+    KafkaProducerHandler, KafkaSource, LogHandler, OffsetTracker, PipelineEnvelope, PipelineExt,
+    PipelineRunner, Stage, StageResult,
+};
+use sentry_arroyo::types::{Topic, TopicOrPartition};
+
+/// A Stage that reverses the string payload.
+struct ReverseStage;
+
+impl Stage for ReverseStage {
+    type In = KafkaPayload;
+    type Out = KafkaPayload;
+
+    async fn process(&self, envelope: PipelineEnvelope<KafkaPayload>) -> StageResult<KafkaPayload> {
+        let reversed = envelope.map_payload(|p| {
+            let bytes = p.payload().unwrap();
+            let s = std::str::from_utf8(bytes).unwrap();
+            let reversed: String = s.chars().rev().collect();
+            println!("transforming: {:?} -> {:?}", s, reversed);
+            KafkaPayload::new(
+                p.key().cloned(),
+                p.headers().cloned(),
+                Some(reversed.into_bytes()),
+            )
+        });
+
+        StageResult::Emit(reversed)
+    }
+
+    fn name(&self) -> &'static str {
+        "reverse_string"
+    }
+}
+
+#[tokio::main]
+async fn main() {
+    tracing_subscriber::fmt::init();
+
+    let consumer_config = KafkaConfig::new_consumer_config(
+        vec!["0.0.0.0:9092".to_string()],
+        "my_group".to_string(),
+        InitialOffset::Latest,
+        false,
+        30_000,
+        None,
+    );
+    let source = KafkaSource::new(consumer_config, &[Topic::new("test_in")]);
+
+    let producer_config = KafkaConfig::new_producer_config(vec!["0.0.0.0:9092".to_string()], None);
+    // The pipeline reads left to right, top to bottom:
+    //   stream                     — async stream of Kafka messages
+    //     .apply(&stage)           — transform/filter/batch each message
+    //     .on_next(&handler)       — side-effect on successful items (produce, upload)
+    //     .on_reject(&handler)     — handle rejected items (DLQ, log)
+    //     .commit(&mut tracker)    — track offsets, flush on interval
+    //
+    // PipelineRunner::run() handles the lifecycle:
+    //   - calls the closure once per partition assignment
+    //   - on rebalance: closure is called again with a fresh stream
+    //   - on shutdown or stream end: exits
+    let result = PipelineRunner::run(&source, |stream, committer| async {
+        let reverse = ReverseStage;
+        let producer = KafkaProducer::new(producer_config.clone());
+        let produce_handler =
+            KafkaProducerHandler::new(producer, TopicOrPartition::Topic(Topic::new("test_out")));
+        let error_handler = LogHandler;
+        let mut tracker = OffsetTracker::new(Duration::from_secs(1), committer);
+
+        stream
+            .apply(&reverse)
+            .on_next(&produce_handler)
+            .on_reject(&error_handler)
+            .commit(&mut tracker)
+            .await
+    })
+    .await;
+
+    if let Err(e) = result {
+        tracing::error!("Pipeline stopped: {}", e);
+    }
+}

rust-arroyo/src/backends/kafka/mod.rs

@@ -82,7 +82,7 @@ impl KafkaConsumerState {
 }
 
 /// Treat recoverable `librdkafka` errors as an empty poll the same way Python treats `KafkaError._TRANSPORT`.
-fn kafka_poll_error_is_recoverable(err: &KafkaError) -> bool {
+pub fn kafka_poll_error_is_recoverable(err: &KafkaError) -> bool {
     matches!(
         err,
         KafkaError::MessageConsumption(RDKafkaErrorCode::BrokerTransportFailure)

rust-arroyo/src/processing/mod.rs

@@ -22,6 +22,7 @@ use crate::{counter, timer};
 pub mod dlq;
 mod metrics_buffer;
 pub mod strategies;
+pub mod stream;
 
 use strategies::{ProcessingStrategy, ProcessingStrategyFactory};
 

rust-arroyo/src/processing/stream/ext.rs

@@ -0,0 +1,199 @@
+use std::time::Instant;
+
+use futures::stream::Stream;
+use futures::StreamExt;
+
+use super::offset_tracker::OffsetTracker;
+use crate::{counter, timer};
+
+use super::handlers::next::NextHandler;
+use super::handlers::rejection::{RejectionHandler, RejectionMetadata};
+use super::pipeline_envelope::PipelineEnvelope;
+use super::stage::{PipelineExit, Stage, StageResult};
+
+/// Run a stage on an envelope and record metrics.
+async fn run_stage<S: Stage>(stage: &S, envelope: PipelineEnvelope<S::In>) -> StageResult<S::Out> {
+    let start = Instant::now();
+    let result = stage.process(envelope).await;
+
+    timer!("arroyo.stage.duration", start.elapsed(), "stage" => stage.name());
+    match &result {
+        StageResult::Emit(_) => {
+            counter!("arroyo.stage.success", 1, "stage" => stage.name());
+        }
+        StageResult::Drop { .. } => {
+            counter!("arroyo.stage.drop", 1, "stage" => stage.name());
+        }
+        StageResult::Skip => {
+            counter!("arroyo.stage.skip", 1, "stage" => stage.name());
+        }
+        StageResult::Reject { .. } => {
+            counter!("arroyo.stage.reject", 1, "stage" => stage.name());
+        }
+        StageResult::Fail(_) => {
+            counter!("arroyo.stage.fail", 1, "stage" => stage.name());
+        }
+        StageResult::Exit(_) => {}
+    }
+
+    result
+}
+
+/// Extension trait that adds pipeline combinators to any
+/// Stream<Item = StageResult<T>>.
+///
+/// Combinators: `.apply()`, `.apply_concurrent()`, `.on_next()`,
+/// `.on_reject()`, `.commit()`.
+///
+/// See `PipelineRunner` for the recommended way to run a pipeline
+/// with rebalance handling.
+pub trait PipelineExt<T: Send>: Stream<Item = StageResult<T>> + Sized {
+    /// Apply a processing stage sequentially to each Emit envelope.
+    /// Equivalent to apply_concurrent(stage, 1).
+    fn apply<'a, S>(self, stage: &'a S) -> impl Stream<Item = StageResult<S::Out>> + 'a
+    where
+        S: Stage<In = T>,
+        Self: 'a,
+        T: 'a,
+    {
+        self.apply_concurrent(stage, 1)
+    }
+
+    /// Apply a processing stage concurrently to up to `concurrency` Emit
+    /// envelopes at once. Results are yielded in input order.
+    /// Non-Emit items pass through immediately.
+    fn apply_concurrent<'a, S>(
+        self,
+        stage: &'a S,
+        concurrency: usize,
+    ) -> impl Stream<Item = StageResult<S::Out>> + 'a
+    where
+        S: Stage<In = T>,
+        Self: 'a,
+        T: 'a,
+    {
+        assert!(concurrency > 0, "concurrency must be at least 1");
+        self.map(move |item| async move {
+            match item {
+                StageResult::Emit(e) => run_stage(stage, e).await,
+                StageResult::Drop { metadata } => StageResult::Drop { metadata },
+                StageResult::Skip => StageResult::Skip,
+                StageResult::Reject {
+                    metadata,
+                    raw,
+                    reason,
+                } => StageResult::Reject {
+                    metadata,
+                    raw,
+                    reason,
+                },
+                StageResult::Fail(err) => StageResult::Fail(err),
+                StageResult::Exit(reason) => StageResult::Exit(reason),
+            }
+        })
+        .buffered(concurrency)
+    }
+
+    /// Call the next handler for each Emit envelope.
+    /// If the handler fails, the item becomes Fail.
+    /// All other variants pass through untouched.
+    fn on_next<'a, H>(self, handler: &'a H) -> impl Stream<Item = StageResult<T>> + 'a
+    where
+        H: NextHandler<T>,
+        Self: 'a,
+        T: 'a,
+    {
+        self.then(move |item| async move {
+            let envelope = match item {
+                StageResult::Emit(e) => e,
+                other => return other,
+            };
+
+            match handler.handle(&envelope).await {
+                Ok(()) => StageResult::Emit(envelope),
+                Err(produce_err) => StageResult::Fail(produce_err),
+            }
+        })
+    }
+
+    /// Call the rejection handler for each Reject item.
+    /// All other variants pass through untouched.
+    /// If the handler itself fails, the item becomes Fail.
+    fn on_reject<'a, H>(self, handler: &'a H) -> impl Stream<Item = StageResult<T>> + 'a
+    where
+        H: RejectionHandler,
+        Self: 'a,
+        T: 'a,
+    {
+        self.then(move |item| async move {
+            match item {
+                StageResult::Reject {
+                    metadata,
+                    raw,
+                    reason,
+                } => {
+                    let rejected = RejectionMetadata {
+                        metadata: metadata.clone(),
+                        raw: raw.clone(),
+                        reason,
+                    };
+
+                    match handler.handle(&rejected).await {
+                        Ok(()) => StageResult::Reject {
+                            metadata,
+                            raw,
+                            reason,
+                        },
+                        Err(handler_err) => StageResult::Fail(handler_err),
+                    }
+                }
+                other => other,
+            }
+        })
+    }
+
+    /// Terminal: drive the pipeline to completion.
+    /// Tracks offsets for Emit, Drop, and Reject items.
+    /// Returns the exit reason (Rebalance, Shutdown, or Complete).
+    /// Fail stops the pipeline with an error.
+    #[allow(async_fn_in_trait)]
+    async fn commit(
+        self,
+        tracker: &mut OffsetTracker<'_>,
+    ) -> Result<PipelineExit, Box<dyn std::error::Error + Send>> {
+        let mut stream = Box::pin(self);
+
+        while let Some(item) = stream.next().await {
+            match item {
+                StageResult::Emit(envelope) => {
+                    tracker.track(envelope.metadata.partition, envelope.metadata.offset + 1);
+                    tracker.record_latency(envelope.metadata.timestamp);
+                }
+                StageResult::Drop { metadata } => {
+                    tracker.track(metadata.partition, metadata.offset + 1);
+                }
+                StageResult::Skip => {}
+                StageResult::Reject { metadata, .. } => {
+                    tracker.track(metadata.partition, metadata.offset + 1);
+                }
+                StageResult::Fail(err) => {
+                    let _ = tracker.flush();
+                    return Err(err);
+                }
+                StageResult::Exit(reason) => {
+                    tracker.flush()?;
+                    return Ok(reason);
+                }
+            }
+
+            tracker.maybe_commit()?;
+        }
+
+        // Stream ended naturally (no Exit item)
+        tracker.flush()?;
+        Ok(PipelineExit::Complete)
+    }
+}
+
+// Blanket impl: any Stream of StageResult<T> gets these methods.
+impl<T: Send, S> PipelineExt<T> for S where S: Stream<Item = StageResult<T>> {}