laminar_core/streaming/source/
mod.rs

1//! Source — entry point for data into a streaming pipeline.
2
3use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
4use std::sync::{Arc, OnceLock};
5use std::time::Duration;
6
7use arrow::array::RecordBatch;
8use arrow::datatypes::SchemaRef;
9use tokio::sync::Semaphore;
10
11use super::arrow_admission::QueuedArrowBatch;
12use super::channel::{channel_with_config, AsyncConsumer, Producer};
13use super::config::SourceConfig;
14use super::error::{StreamingError, TryPushError};
15use super::sink::Sink;
16
17/// Trait for types that can be streamed through a Source.
18pub trait Record: Clone + Send + Sized + 'static {
19    /// Returns the Arrow schema for this record type.
20    fn schema() -> SchemaRef;
21
22    /// Converts this record to an Arrow `RecordBatch`.
23    ///
24    /// The batch will contain a single row with this record's data.
25    fn to_record_batch(&self) -> RecordBatch;
26
27    /// Returns the event time for this record, if applicable.
28    ///
29    /// Event time is used for watermark generation and window assignment.
30    /// Returns `None` if the record doesn't have an event time.
31    fn event_time(&self) -> Option<i64> {
32        None
33    }
34
35    /// Converts a batch of records to an Arrow `RecordBatch`.
36    ///
37    /// The default implementation converts each record individually and concatenates them.
38    /// Derived implementations can override this to optimize allocation and copying.
39    fn to_record_batch_from_iter<I>(records: I) -> RecordBatch
40    where
41        I: IntoIterator<Item = Self>,
42    {
43        let batches: Vec<RecordBatch> = records.into_iter().map(|r| r.to_record_batch()).collect();
44        if batches.is_empty() {
45            return RecordBatch::new_empty(Self::schema());
46        }
47        arrow::compute::concat_batches(&Self::schema(), &batches)
48            .unwrap_or_else(|_| RecordBatch::new_empty(Self::schema()))
49    }
50}
51
52/// Internal message type that wraps records and control signals.
53pub(crate) enum SourceMessage<T> {
54    /// A data record.
55    Record(T),
56
57    /// A batch of Arrow records.
58    Batch(QueuedArrowBatch),
59}
60
61impl<T: Record> SourceMessage<T> {
62    pub(crate) fn into_batch(self) -> RecordBatch {
63        match self {
64            Self::Record(record) => record.to_record_batch(),
65            Self::Batch(batch) => batch.into_batch(),
66        }
67    }
68}
69
70/// Shared state for watermark tracking.
71struct SourceWatermark {
72    /// Current watermark value.
73    /// Atomically updated to support multi-producer scenarios.
74    /// Wrapped in `Arc` so the checkpoint manager can read it without locking.
75    current: Arc<AtomicI64>,
76}
77
78impl SourceWatermark {
79    fn new() -> Self {
80        Self {
81            current: Arc::new(AtomicI64::new(i64::MIN)),
82        }
83    }
84
85    fn from_arc(arc: Arc<AtomicI64>) -> Self {
86        Self { current: arc }
87    }
88
89    fn update(&self, timestamp: i64) {
90        // Only advance watermark, never go backwards
91        let mut current = self.current.load(Ordering::Acquire);
92        while timestamp > current {
93            match self.current.compare_exchange_weak(
94                current,
95                timestamp,
96                Ordering::AcqRel,
97                Ordering::Acquire,
98            ) {
99                Ok(_) => break,
100                Err(actual) => current = actual,
101            }
102        }
103    }
104
105    fn restore_for_recovery(&self, timestamp: i64) {
106        self.current.store(timestamp, Ordering::Release);
107    }
108
109    fn get(&self) -> i64 {
110        self.current.load(Ordering::Acquire)
111    }
112
113    fn arc(&self) -> Arc<AtomicI64> {
114        Arc::clone(&self.current)
115    }
116}
117
118/// Shared state for a Source/Sink pair.
119struct SourceInner<T: Record> {
120    /// Channel producer for sending records.
121    producer: Producer<SourceMessage<T>>,
122
123    arrow_budget: Arc<Semaphore>,
124    max_queued_bytes: usize,
125
126    /// Watermark state.
127    watermark: SourceWatermark,
128
129    /// Schema for type validation.
130    schema: SchemaRef,
131
132    /// Source name (for debugging/metrics).
133    name: Option<String>,
134
135    /// Monotonic sequence counter, incremented on each successful push.
136    /// Wrapped in `Arc` so the checkpoint manager can read it without locking.
137    sequence: Arc<AtomicU64>,
138
139    /// Event-time column name set via programmatic API.
140    /// Read once at pipeline startup, not on the hot path.
141    event_time_column: OnceLock<String>,
142
143    /// Max out-of-orderness bound, paired with `event_time_column`.
144    /// Read once at pipeline startup, not on the hot path.
145    max_out_of_orderness: OnceLock<Duration>,
146}
147
148/// A streaming data source. Cloneable for multi-producer use.
149///
150/// Arrow pushes share one byte budget, held by each batch until the sink receives it.
151/// Generic record pushes are count bounded only; arbitrary `T` heap storage is not measured.
152pub struct Source<T: Record> {
153    inner: Arc<SourceInner<T>>,
154}
155
156impl<T: Record> Source<T> {
157    /// Creates a new Source/Sink pair.
158    pub(crate) fn new(config: SourceConfig) -> (Self, Sink<T>) {
159        let (source, consumer) = Self::with_consumer(config);
160        let sink = Sink::new(consumer, source.schema());
161        (source, sink)
162    }
163
164    fn with_consumer(config: SourceConfig) -> (Self, AsyncConsumer<SourceMessage<T>>) {
165        let channel_config = config.channel;
166        let (producer, consumer) = channel_with_config::<SourceMessage<T>>(&channel_config);
167
168        let schema = T::schema();
169
170        let inner = Arc::new(SourceInner {
171            producer,
172            // Constructors remain infallible. Invalid byte limits fail Arrow admission,
173            // without passing an unsupported permit count to Tokio's constructor.
174            arrow_budget: Arc::new(Semaphore::new(
175                config.max_queued_bytes.min(super::MAX_SOURCE_QUEUED_BYTES),
176            )),
177            max_queued_bytes: config.max_queued_bytes,
178            watermark: SourceWatermark::new(),
179            schema,
180            name: config.name,
181            sequence: Arc::new(AtomicU64::new(0)),
182            event_time_column: OnceLock::new(),
183            max_out_of_orderness: OnceLock::new(),
184        });
185
186        (Self { inner }, consumer)
187    }
188
189    /// Pushes a record. Non-blocking — returns `ChannelFull` if the buffer is full.
190    ///
191    /// # Errors
192    ///
193    /// Returns `StreamingError::ChannelFull` if the count buffer is full, or
194    /// `StreamingError::Disconnected` if the sink has been dropped.
195    pub fn push(&self, record: T) -> Result<(), StreamingError> {
196        self.try_push(record).map_err(|error| error.error)
197    }
198
199    /// Pushes a record, returning it on failure.
200    ///
201    /// # Errors
202    ///
203    /// Returns `TryPushError` containing the record if the channel is full or disconnected.
204    pub fn try_push(&self, record: T) -> Result<(), TryPushError<T>> {
205        let event_time = record.event_time();
206
207        self.inner
208            .producer
209            .try_push(SourceMessage::Record(record))
210            .map_err(|error| match error.value {
211                SourceMessage::Record(r) => TryPushError {
212                    value: r,
213                    error: error.error,
214                },
215                SourceMessage::Batch(_) => unreachable!("only Record is pushed here"),
216            })?;
217
218        if let Some(event_time) = event_time {
219            self.inner.watermark.update(event_time);
220        }
221        self.inner.sequence.fetch_add(1, Ordering::Relaxed);
222        Ok(())
223    }
224
225    /// Pushes multiple records (cloned). Stops at the first failure.
226    pub fn push_batch(&self, records: &[T]) -> usize
227    where
228        T: Clone,
229    {
230        self.push_batch_drain(records.iter().cloned())
231    }
232
233    /// Pushes records from an iterator, consuming them (zero-clone).
234    /// Stops at the first failure. Returns the number pushed.
235    pub fn push_batch_drain<I>(&self, records: I) -> usize
236    where
237        I: IntoIterator<Item = T>,
238    {
239        let mut count = 0;
240        for record in records {
241            if self.push(record).is_err() {
242                break;
243            }
244            count += 1;
245        }
246        count
247    }
248
249    /// Pushes an Arrow `RecordBatch` directly.
250    ///
251    /// This is more efficient than pushing individual records when you
252    /// already have data in Arrow format.
253    /// The nonblocking reservation covers retained Arrow storage until the sink receives the
254    /// batch, and releases on failed admission. Received batches and downstream retention have
255    /// separate owners.
256    ///
257    /// # Errors
258    ///
259    /// Returns `StreamingError::SchemaMismatch` if the batch schema doesn't match.
260    /// Returns `StreamingError::BatchTooLarge` if this batch exceeds the whole byte budget,
261    /// `ChannelFull` on count/byte saturation, `Disconnected` if the sink has been dropped,
262    /// or `InvalidConfig` for an invalid byte limit. Failures do not advance the sequence.
263    pub fn push_arrow(&self, batch: RecordBatch) -> Result<(), StreamingError> {
264        // Validate schema matches (skip for type-erased sources with empty schema)
265        if !self.inner.schema.fields().is_empty() && batch.schema() != self.inner.schema {
266            return Err(StreamingError::SchemaMismatch {
267                expected: self
268                    .inner
269                    .schema
270                    .fields()
271                    .iter()
272                    .map(|f| f.name().clone())
273                    .collect(),
274                actual: batch
275                    .schema()
276                    .fields()
277                    .iter()
278                    .map(|f| f.name().clone())
279                    .collect(),
280            });
281        }
282
283        if self.is_closed() {
284            return Err(StreamingError::Disconnected);
285        }
286        let queued =
287            QueuedArrowBatch::admit(batch, &self.inner.arrow_budget, self.inner.max_queued_bytes)?;
288        self.inner
289            .producer
290            .try_push(SourceMessage::Batch(queued))
291            .map_err(|error| error.error)?;
292
293        self.inner.sequence.fetch_add(1, Ordering::Relaxed);
294        Ok(())
295    }
296
297    /// Emits a watermark timestamp.
298    ///
299    /// Watermarks signal that no events with timestamps less than or equal
300    /// to this value will arrive in the future. This enables window triggers
301    /// and garbage collection.
302    ///
303    /// Watermarks are monotonically increasing - if a lower timestamp is
304    /// passed, it will be ignored.
305    pub fn watermark(&self, timestamp: i64) {
306        // The shared atomic is the authoritative watermark: the pipeline's
307        // watermark UDF, late-row filter, and checkpoint registration all
308        // read it via `watermark_atomic()`. The sink receives data only,
309        // so there is no in-band watermark message to emit.
310        self.inner.watermark.update(timestamp);
311    }
312
313    /// Replaces the shared watermark with an exact recovered value.
314    ///
315    /// Unlike [`Self::watermark`], this may lower the watermark. The caller
316    /// must hold the intake/compute recovery fence and ensure every clone and
317    /// producer of this source is quiesced until restoration completes. Use
318    /// [`Self::watermark`] during normal operation.
319    pub fn restore_watermark_for_recovery(&self, timestamp: i64) {
320        self.inner.watermark.restore_for_recovery(timestamp);
321    }
322
323    /// Returns the current watermark value.
324    #[must_use]
325    pub fn current_watermark(&self) -> i64 {
326        self.inner.watermark.get()
327    }
328
329    /// Returns the schema for this source.
330    #[must_use]
331    pub fn schema(&self) -> SchemaRef {
332        Arc::clone(&self.inner.schema)
333    }
334
335    /// Returns the source name, if configured.
336    #[must_use]
337    pub fn name(&self) -> Option<&str> {
338        self.inner.name.as_deref()
339    }
340
341    /// Returns true once the sink has been dropped.
342    #[must_use]
343    pub fn is_closed(&self) -> bool {
344        self.inner.producer.is_closed()
345    }
346
347    /// Returns the number of pending items in the buffer.
348    #[must_use]
349    pub fn pending(&self) -> usize {
350        self.inner.producer.len()
351    }
352
353    /// Returns the buffer capacity.
354    #[must_use]
355    pub fn capacity(&self) -> usize {
356        self.inner.producer.capacity()
357    }
358
359    /// Charged Arrow bytes queued for the sink, across all producers.
360    /// Generic records and batches the sink has received are excluded.
361    #[must_use]
362    pub fn queued_arrow_bytes(&self) -> usize {
363        self.inner
364            .max_queued_bytes
365            .min(super::MAX_SOURCE_QUEUED_BYTES)
366            - self.inner.arrow_budget.available_permits()
367    }
368
369    /// Configured Arrow-byte budget for this source.
370    #[must_use]
371    pub fn max_queued_bytes(&self) -> usize {
372        self.inner.max_queued_bytes
373    }
374
375    /// Returns the current sequence number (total successful pushes).
376    #[must_use]
377    pub fn sequence(&self) -> u64 {
378        self.inner.sequence.load(Ordering::Acquire)
379    }
380
381    /// Returns the shared sequence counter for checkpoint registration.
382    #[must_use]
383    pub fn sequence_counter(&self) -> Arc<AtomicU64> {
384        Arc::clone(&self.inner.sequence)
385    }
386
387    /// Returns the shared watermark atomic for checkpoint registration.
388    #[must_use]
389    pub fn watermark_atomic(&self) -> Arc<AtomicI64> {
390        self.inner.watermark.arc()
391    }
392
393    /// Declare which column in the source data represents event time.
394    ///
395    /// When set, `source.watermark()` enables late-row filtering
396    /// without a SQL `WATERMARK FOR` clause.
397    ///
398    /// Only the first call takes effect; subsequent calls are silently ignored.
399    pub fn set_event_time_column(&self, column: &str) {
400        let _ = self.inner.event_time_column.set(column.to_owned());
401    }
402
403    /// Returns the configured event-time column, if any.
404    #[must_use]
405    pub fn event_time_column(&self) -> Option<String> {
406        self.inner.event_time_column.get().cloned()
407    }
408
409    /// Set the max out-of-orderness bound for watermark generation.
410    ///
411    /// Only the first call takes effect; subsequent calls are silently ignored.
412    pub fn set_max_out_of_orderness(&self, dur: Duration) {
413        let _ = self.inner.max_out_of_orderness.set(dur);
414    }
415
416    /// Returns the configured max out-of-orderness, if any.
417    #[must_use]
418    pub fn max_out_of_orderness(&self) -> Option<Duration> {
419        self.inner.max_out_of_orderness.get().copied()
420    }
421}
422
423impl<T: Record> Clone for Source<T> {
424    fn clone(&self) -> Self {
425        let producer = self.inner.producer.clone();
426        let event_time_col = self.inner.event_time_column.get().cloned();
427        let event_time_column = OnceLock::new();
428        if let Some(col) = event_time_col {
429            let _ = event_time_column.set(col);
430        }
431        let max_ooo = self.inner.max_out_of_orderness.get().copied();
432        let max_out_of_orderness = OnceLock::new();
433        if let Some(dur) = max_ooo {
434            let _ = max_out_of_orderness.set(dur);
435        }
436        Self {
437            inner: Arc::new(SourceInner {
438                producer,
439                arrow_budget: Arc::clone(&self.inner.arrow_budget),
440                max_queued_bytes: self.inner.max_queued_bytes,
441                watermark: SourceWatermark::from_arc(self.inner.watermark.arc()),
442                schema: Arc::clone(&self.inner.schema),
443                name: self.inner.name.clone(),
444                sequence: Arc::clone(&self.inner.sequence),
445                event_time_column,
446                max_out_of_orderness,
447            }),
448        }
449    }
450}
451
452impl<T: Record + std::fmt::Debug> std::fmt::Debug for Source<T> {
453    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
454        f.debug_struct("Source")
455            .field("name", &self.inner.name)
456            .field("pending", &self.pending())
457            .field("capacity", &self.capacity())
458            .field("watermark", &self.current_watermark())
459            .finish()
460    }
461}
462
463/// Creates a new Source/Sink pair with the given buffer size.
464#[must_use]
465pub fn create<T: Record>(buffer_size: usize) -> (Source<T>, Sink<T>) {
466    Source::new(SourceConfig::with_buffer_size(buffer_size))
467}
468
469/// Creates a new Source/Sink pair with custom configuration.
470#[must_use]
471pub fn create_with_config<T: Record>(config: SourceConfig) -> (Source<T>, Sink<T>) {
472    Source::new(config)
473}
474
475/// Creates a schema-planning source without a sink.
476///
477/// The source is disconnected immediately and rejects all input. This is only for private,
478/// unstarted catalogs used to compile schemas; it cannot execute or publish streaming data.
479#[cfg(feature = "cluster")]
480#[must_use]
481pub fn create_for_schema_planning<T: Record>(config: SourceConfig) -> Source<T> {
482    Source::with_consumer(config).0
483}
484
485#[cfg(test)]
486mod tests;
487
488#[cfg(test)]
489mod admission_tests;
laminar_core/streaming/source/mod.rs

laminar_core/streaming/source/
mod.rs
