Add bytes-based flushing to source collectors

Sources read data in and collect it up into batches. Batches are flushed
based on configurable criteria: until now this has been the number of
records collected and how long they have been waiting.

This commit adds another criteria for flushing, based on the memory
usage of the records. `source-batch-max-bytes`, disabled by default,
allows triggering flushing when a certain amount of data has been
collected.

This new feature enables two useful behaviours:

1) Large records: reduce the risk of OOMs when memory-constrained
2) Small records: create larger, more-efficient batches

Memory usage is estimated based on input bytes, rather than by
instrumenting Arrow buffer memory.

Ideally we would implement this based on Arrow buffer usage. However,
Arrow's `ArrayBuilder` trait doesn't expose memory size and we would
need to maintain a lot of code ourselves to make it work. I think it's
reasonable to use input bytes as an approximation.

Author: 46bit

crates/arroyo-connectors/src/nexmark/operator.rs

@@ -281,7 +281,9 @@ impl SourceOperator for NexmarkSourceFunc {
             next_event.bid.as_ref().write_into(&mut bid_builder);
             timestamp_builder.append_value(to_nanos(next_event.event_timetamp) as i64);
 
-            if should_flush(records, flush_time) {
+            // Nexmark generates events synthetically with no raw input bytes.
+            // Pass 0 for bytes buffered.
+            if should_flush(records, 0, flush_time) {
                 collector
                     .collect(RecordBatch::try_new(
                         ctx.out_schema.schema.clone(),

crates/arroyo-formats/src/de.rs

@@ -36,6 +36,9 @@ pub enum FieldValueType<'a> {
 
 struct ContextBuffer {
     buffer: Vec<Box<dyn ArrayBuilder>>,
+    /// Accumulated raw input bytes since last flush.
+    /// Currently measured as pre-deserialization input size rather than Arrow buffer memory.
+    buffered_bytes: usize,
     created: Instant,
 }
 
@@ -49,6 +52,7 @@ impl ContextBuffer {
 
         Self {
             buffer,
+            buffered_bytes: 0,
             created: Instant::now(),
         }
     }
@@ -58,10 +62,11 @@ impl ContextBuffer {
     }
 
     pub fn should_flush(&self) -> bool {
-        should_flush(self.size(), self.created)
+        should_flush(self.size(), self.buffered_bytes, self.created)
     }
 
     pub fn finish(&mut self) -> Vec<ArrayRef> {
+        self.buffered_bytes = 0;
         self.buffer.iter_mut().map(|a| a.finish()).collect()
     }
 }
@@ -170,6 +175,9 @@ enum BufferDecoder {
     JsonDecoder {
         decoder: arrow::json::reader::Decoder,
         buffered_count: usize,
+        /// Accumulated raw input bytes since last flush.
+        /// Currently measured as pre-deserialization input size rather than Arrow buffer memory.
+        buffered_bytes: usize,
         buffered_since: Instant,
     },
 }
@@ -180,9 +188,10 @@ impl BufferDecoder {
             BufferDecoder::Buffer(b) => b.should_flush(),
             BufferDecoder::JsonDecoder {
                 buffered_count,
+                buffered_bytes,
                 buffered_since,
                 ..
-            } => should_flush(*buffered_count, *buffered_since),
+            } => should_flush(*buffered_count, *buffered_bytes, *buffered_since),
         }
     }
 
@@ -203,9 +212,11 @@ impl BufferDecoder {
                 decoder,
                 buffered_since,
                 buffered_count,
+                buffered_bytes,
             } => {
                 *buffered_since = Instant::now();
                 *buffered_count = 0;
+                *buffered_bytes = 0;
                 Some(match bad_data {
                     BadData::Fail { .. } => decoder
                         .flush()
@@ -307,6 +318,13 @@ impl BufferDecoder {
             }
         }
     }
+
+    fn add_input_bytes(&mut self, n: usize) {
+        match self {
+            BufferDecoder::Buffer(b) => b.buffered_bytes += n,
+            BufferDecoder::JsonDecoder { buffered_bytes, .. } => *buffered_bytes += n,
+        }
+    }
 }
 
 pub struct ArrowDeserializer {
@@ -440,6 +458,7 @@ impl ArrowDeserializer {
                     .build_decoder()
                     .unwrap(),
                 buffered_count: 0,
+                buffered_bytes: 0,
                 buffered_since: Instant::now(),
             },
             _ => BufferDecoder::Buffer(ContextBuffer::new(schema_without_additional.clone())),
@@ -497,18 +516,26 @@ impl ArrowDeserializer {
         additional_fields: Option<&HashMap<&str, FieldValueType<'_>>>,
     ) -> Vec<DataflowError> {
         let (count, errors) = match &*self.format {
-            Format::Avro(_) => self.deserialize_slice_avro(msg).await,
+            Format::Avro(_) => {
+                let (count, errors) = self.deserialize_slice_avro(msg).await;
+                if count > 0 {
+                    self.buffer_decoder.add_input_bytes(msg.len());
+                }
+                (count, errors)
+            }
             _ => {
                 let mut count = 0;
-                let errors = FramingIterator::new(self.framing.clone(), msg)
-                    .map(|t| self.deserialize_single(t))
-                    .filter_map(|t| {
-                        if t.is_ok() {
+                let mut errors = vec![];
+                for t in FramingIterator::new(self.framing.clone(), msg) {
+                    let len = t.len();
+                    match self.deserialize_single(t) {
+                        Ok(()) => {
                             count += 1;
+                            self.buffer_decoder.add_input_bytes(len);
                         }
-                        t.err()
-                    })
-                    .collect();
+                        Err(e) => errors.push(e),
+                    }
+                }
                 (count, errors)
             }
         };

crates/arroyo-formats/src/lib.rs

@@ -9,10 +9,16 @@ pub mod de;
 pub mod proto;
 pub mod ser;
 
-pub fn should_flush(size: usize, time: Instant) -> bool {
-    size > 0
-        && (size >= config().pipeline.source_batch_size
-            || time.elapsed() >= *config().pipeline.source_batch_linger)
+pub fn should_flush(buffered_count: usize, buffered_bytes: usize, time: Instant) -> bool {
+    if buffered_count == 0 {
+        return false;
+    }
+    let cfg = &config().pipeline;
+    buffered_count >= cfg.source_batch_size
+        || time.elapsed() >= *cfg.source_batch_linger
+        || cfg
+            .source_batch_max_bytes
+            .is_some_and(|max| buffered_bytes >= max)
 }
 
 pub(crate) fn float_to_json(f: f64) -> Value {
@@ -30,3 +36,53 @@ pub(crate) fn float_to_json(f: f64) -> Value {
         ),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arroyo_rpc::config::update;
+    use std::time::Duration;
+
+    /// Set all flush-related config fields so each test is self-contained
+    /// even if tests run in parallel.
+    fn set_flush_config(batch_size: usize, linger_ms: u64, max_bytes: Option<usize>) {
+        // config() auto-initialises from defaults; update() panics without it
+        let _ = config();
+        update(|c| {
+            c.pipeline.source_batch_size = batch_size;
+            c.pipeline.source_batch_linger = Duration::from_millis(linger_ms).into();
+            c.pipeline.source_batch_max_bytes = max_bytes;
+        });
+    }
+
+    #[test]
+    fn bytes_threshold_triggers_flush() {
+        set_flush_config(512, 10_000, Some(100));
+        assert!(should_flush(1, 100, Instant::now()));
+    }
+
+    #[test]
+    fn below_bytes_threshold_no_flush() {
+        set_flush_config(512, 10_000, Some(100));
+        assert!(!should_flush(1, 99, Instant::now()));
+    }
+
+    #[test]
+    fn bytes_threshold_disabled() {
+        set_flush_config(512, 10_000, None);
+        assert!(!should_flush(1, 1_000_000, Instant::now()));
+        assert!(!should_flush(1, usize::MAX, Instant::now()));
+    }
+
+    #[test]
+    fn count_threshold_triggers_flush() {
+        set_flush_config(10, 10_000, None);
+        assert!(should_flush(10, 0, Instant::now()));
+    }
+
+    #[test]
+    fn linger_triggers_flush() {
+        set_flush_config(512, 100, None);
+        assert!(should_flush(1, 0, Instant::now() - Duration::from_secs(1)));
+    }
+}

crates/arroyo-rpc/default.toml

@@ -5,6 +5,7 @@ job-controller = "controller"
 [pipeline]
 source-batch-size = 512
 source-batch-linger = "100ms"
+# source-batch-max-bytes = 10485760  # flush after ~10MB of input; disabled by default
 update-aggregate-flush-interval = "1s"
 allowed-restarts = 20
 worker-heartbeat-timeout = "30s"

crates/arroyo-rpc/src/config.rs

@@ -519,6 +519,10 @@ pub struct PipelineConfig {
     /// Batch linger time (how long to wait before flushing)
     pub source_batch_linger: HumanReadableDuration,
 
+    /// Maximum buffered bytes before flushing (None = disabled).
+    /// Currently measured as raw input bytes.
+    pub source_batch_max_bytes: Option<usize>,
+
     /// How often to flush aggregates
     pub update_aggregate_flush_interval: HumanReadableDuration,