feat(runtime): Implement pull based arroyo (hackweek)

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,104 @@
+/// Pull-based version of transform_and_produce.
+///
+/// Pipeline:
+///   KafkaSource → apply(reverse) → on_next(produce) → on_reject(log) → commit
+///
+/// `PipelineRunner::run_pipeline()` handles the rebalance restart loop —
+/// the build closure is called once per partition assignment with fresh
+/// stages and handlers.
+extern crate sentry_arroyo;
+
+use std::time::Duration;
+
+use futures::Stream;
+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, Pipeline, 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"
+    }
+}
+
+struct TransformAndProducePipeline {
+    reverse: ReverseStage,
+    produce_handler: KafkaProducerHandler,
+    error_handler: LogHandler,
+}
+
+impl Pipeline for TransformAndProducePipeline {
+    type Output = KafkaPayload;
+
+    fn stream(
+        self,
+        source: impl Stream<Item = StageResult<KafkaPayload>> + Send,
+    ) -> impl Stream<Item = StageResult<KafkaPayload>> + Send {
+        source
+            .apply(self.reverse)
+            .on_next(self.produce_handler)
+            .on_reject(self.error_handler)
+    }
+}
+
+#[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);
+
+    let result = PipelineRunner::run(&source, Duration::from_secs(1), || {
+        let producer = KafkaProducer::new(producer_config.clone()).unwrap();
+        TransformAndProducePipeline {
+            reverse: ReverseStage,
+            produce_handler: KafkaProducerHandler::new(
+                producer,
+                TopicOrPartition::Topic(Topic::new("test_out")),
+            ),
+            error_handler: LogHandler,
+        }
+    })
+    .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/batch/buffer.rs

@@ -0,0 +1,26 @@
+/// Accumulates items for batching. Implementations define how items
+/// are stored, how their byte size is measured, and what the flushed
+/// output type is.
+///
+/// The batch stage calls `push()` for each item, uses the returned
+/// byte count to track size-based thresholds, and calls `flush()`
+/// when a threshold is reached.
+pub trait Buffer<T>: Send + Sync {
+    /// The type returned by `flush()`. For collecting buffers this is `Vec<T>`.
+    /// For merging buffers this can be a single aggregated value.
+    type Output: Send + Sync;
+
+    /// Add an item to the buffer. Returns the item's size in bytes.
+    fn push(&mut self, item: T) -> u64;
+
+    /// Number of items in the buffer.
+    fn len(&self) -> u64;
+
+    /// Whether the buffer is empty.
+    fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// Drain the buffer and return the accumulated output.
+    fn flush(&mut self) -> Self::Output;
+}

rust-arroyo/src/processing/stream/batch/flush_timer.rs

@@ -0,0 +1,139 @@
+use std::time::Duration;
+
+use tokio::time::Interval;
+
+use crate::processing::stream::offset_tracker::ensure_time_updater;
+
+/// A monotonic time source for `FlushTimer`.
+///
+/// Production uses `CoarseClock` (coarsetime's cached global, ~1ns reads).
+/// Tests can inject a controllable clock for deterministic timing.
+pub trait Clock: Send + Sync {
+    fn now(&self) -> coarsetime::Instant;
+}
+
+/// Production clock — reads from coarsetime's cached global.
+pub struct CoarseClock;
+
+impl Clock for CoarseClock {
+    fn now(&self) -> coarsetime::Instant {
+        coarsetime::Instant::recent()
+    }
+}
+
+/// Manages time-based flush triggers for `apply_with_timer`.
+///
+/// Uses coarsetime watermarks checked on each interval tick.
+/// No per-item timer wheel operations.
+///
+/// Two triggers:
+/// - **Idle**: flush if no upstream item arrives within `idle_timeout`.
+///   Watermark `last_activity` resets on every item.
+/// - **Cadence**: flush at `max_cadence` after the first item accumulated.
+///   Watermark `batch_start` set once per batch.
+pub struct FlushTimer<C: Clock> {
+    clock: C,
+    /// Interval tick — poll this in `select!`.
+    pub interval: Interval,
+
+    idle_dur: coarsetime::Duration,
+    cadence_dur: coarsetime::Duration,
+
+    batch_start: coarsetime::Instant,
+    last_activity: coarsetime::Instant,
+}
+
+/// Far-future offset used as a sentinel — makes `is_active()` return false
+/// and `should_flush()` return false without branches.
+fn far_future() -> coarsetime::Duration {
+    coarsetime::Duration::from_secs(365 * 24 * 3600)
+}
+
+impl FlushTimer<CoarseClock> {
+    /// Create a production `FlushTimer` with coarsetime and a tokio interval.
+    pub fn new(idle_timeout: Option<Duration>, max_cadence: Option<Duration>) -> Self {
+        ensure_time_updater();
+
+        let min_dur = [idle_timeout, max_cadence]
+            .into_iter()
+            .flatten()
+            .min()
+            .unwrap_or(Duration::from_secs(1));
+        let tick = (min_dur / 10)
+            .max(Duration::from_millis(1))
+            .min(Duration::from_millis(100));
+
+        Self::with_clock(
+            CoarseClock,
+            tokio::time::interval(tick),
+            idle_timeout,
+            max_cadence,
+        )
+    }
+}
+
+impl<C: Clock> FlushTimer<C> {
+    /// Create a `FlushTimer` with an injected clock and interval.
+    pub fn with_clock(
+        clock: C,
+        interval: Interval,
+        idle_timeout: Option<Duration>,
+        max_cadence: Option<Duration>,
+    ) -> Self {
+        let far = clock.now() + far_future();
+
+        Self {
+            clock,
+            interval,
+            idle_dur: idle_timeout.map_or(far_future(), |d| d.into()),
+            cadence_dur: max_cadence.map_or(far_future(), |d| d.into()),
+            batch_start: far,
+            last_activity: far,
+        }
+    }
+
+    /// Whether a batch is currently accumulating.
+    pub fn is_active(&self) -> bool {
+        self.batch_start <= self.clock.now()
+    }
+
+    /// An item was accumulated. Updates watermarks (~1ns).
+    pub fn on_accumulate(&mut self) {
+        let now = self.clock.now();
+        self.last_activity = now;
+        if self.batch_start > now {
+            self.batch_start = now;
+        }
+    }
+
+    /// The batch was flushed. Unsets watermarks.
+    pub fn on_flush(&mut self) {
+        self.unset();
+    }
+
+    /// Duration since the batch started accumulating.
+    /// Returns zero if no batch is active.
+    pub fn batch_time(&self) -> Duration {
+        let now = self.clock.now();
+        if self.batch_start > now {
+            Duration::ZERO
+        } else {
+            now.duration_since(self.batch_start).into()
+        }
+    }
+
+    /// Check if a flush trigger has fired based on watermarks.
+    /// Called on each interval tick.
+    pub fn should_flush(&self) -> bool {
+        let now = self.clock.now();
+        now.duration_since(self.batch_start) >= self.cadence_dur
+            || now.duration_since(self.last_activity) >= self.idle_dur
+    }
+
+    /// Reset watermarks to far-future sentinels.
+    fn unset(&mut self) {
+        let far = self.clock.now() + far_future();
+        self.batch_start = far;
+        self.last_activity = far;
+    }
+}

rust-arroyo/src/processing/stream/batch/mod.rs

@@ -0,0 +1,4 @@
+pub mod buffer;
+pub mod flush_timer;
+pub mod stage;
+pub mod triggers;