laminar_core/streaming/subscription/
mod.rs

1//! Subscription — receive records from a Sink.
2
3use std::sync::Arc;
4use std::time::Duration;
5
6use arrow::array::RecordBatch;
7use arrow::datatypes::SchemaRef;
8use tokio::sync::broadcast;
9
10use super::error::RecvError;
11use super::source::{Record, SourceMessage};
12
13/// A subscription to a streaming sink. Each subscriber independently receives
14/// every message via broadcast.
15pub struct Subscription<T: Record> {
16    rx: broadcast::Receiver<SourceMessage<T>>,
17    schema: SchemaRef,
18    closed: bool,
19}
20
21impl<T: Record> Subscription<T> {
22    pub(crate) fn new(rx: broadcast::Receiver<SourceMessage<T>>, schema: SchemaRef) -> Self {
23        Self {
24            rx,
25            schema,
26            closed: false,
27        }
28    }
29
30    /// Non-blocking poll. Returns the next batch.
31    /// Returns `None` on empty or closed channel. Check `is_disconnected()`
32    /// to distinguish.
33    pub fn poll(&mut self) -> Option<RecordBatch> {
34        loop {
35            match self.rx.try_recv() {
36                Ok(msg) => return Some(to_batch(msg)),
37                Err(broadcast::error::TryRecvError::Empty) => return None,
38                Err(broadcast::error::TryRecvError::Closed) => {
39                    self.closed = true;
40                    return None;
41                }
42                Err(broadcast::error::TryRecvError::Lagged(_)) => {}
43            }
44        }
45    }
46
47    /// Async receive. Awaits the next batch.
48    ///
49    /// # Errors
50    ///
51    /// Returns `RecvError::Disconnected` if the source has been dropped.
52    pub async fn recv_async(&mut self) -> Result<RecordBatch, RecvError> {
53        loop {
54            match self.rx.recv().await {
55                Ok(msg) => return Ok(to_batch(msg)),
56                Err(broadcast::error::RecvError::Closed) => {
57                    self.closed = true;
58                    return Err(RecvError::Disconnected);
59                }
60                Err(broadcast::error::RecvError::Lagged(_)) => {}
61            }
62        }
63    }
64
65    /// Blocking receive. Uses tokio's waker-based `blocking_recv`.
66    ///
67    /// # Errors
68    ///
69    /// Returns `RecvError::Disconnected` if the source has been dropped.
70    pub fn recv(&mut self) -> Result<RecordBatch, RecvError> {
71        loop {
72            match self.rx.blocking_recv() {
73                Ok(msg) => return Ok(to_batch(msg)),
74                Err(broadcast::error::RecvError::Closed) => {
75                    self.closed = true;
76                    return Err(RecvError::Disconnected);
77                }
78                Err(broadcast::error::RecvError::Lagged(_)) => {}
79            }
80        }
81    }
82
83    /// Blocking receive with timeout. Requires a tokio runtime in the current
84    /// thread context.
85    ///
86    /// # Errors
87    ///
88    /// Returns `RecvError::Timeout` or `RecvError::Disconnected`.
89    pub fn recv_timeout(&mut self, timeout: Duration) -> Result<RecordBatch, RecvError> {
90        let handle = tokio::runtime::Handle::current();
91        match handle.block_on(tokio::time::timeout(timeout, self.recv_async())) {
92            Ok(Ok(batch)) => Ok(batch),
93            Ok(Err(e)) => Err(e),
94            Err(_) => Err(RecvError::Timeout),
95        }
96    }
97
98    /// Returns true if the channel has been observed closed.
99    #[must_use]
100    pub fn is_disconnected(&self) -> bool {
101        self.closed
102    }
103
104    /// Returns the schema for records in this subscription.
105    #[must_use]
106    pub fn schema(&self) -> SchemaRef {
107        Arc::clone(&self.schema)
108    }
109}
110
111fn to_batch<T: Record>(msg: SourceMessage<T>) -> RecordBatch {
112    match msg {
113        SourceMessage::Record(r) => r.to_record_batch(),
114        SourceMessage::Batch(b) => b,
115    }
116}
117
118impl<T: Record + std::fmt::Debug> std::fmt::Debug for Subscription<T> {
119    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120        f.debug_struct("Subscription")
121            .field("closed", &self.closed)
122            .field("schema", &self.schema)
123            .finish_non_exhaustive()
124    }
125}
126
127#[cfg(test)]
128mod tests;
laminar_core/streaming/subscription/mod.rs

laminar_core/streaming/subscription/
mod.rs
