laminar_db/subscription/
portal.rs

1//! Per-subscriber cursor over the shared subscription log.
2
3use std::sync::Arc;
4
5use arrow_array::RecordBatch;
6use arrow_schema::SchemaRef;
7use datafusion::physical_expr::PhysicalExpr;
8use futures::FutureExt;
9
10use super::registry::{ChargedUpdate, MvUpdate, SubscriptionRead, SubscriptionReader};
11
12/// Keeps the process-wide subscription charge alive with an emitted batch.
13#[doc(hidden)]
14#[derive(Clone)]
15pub struct SubscriptionFrameLease {
16    _owner: ChargedUpdate,
17}
18
19impl std::fmt::Debug for SubscriptionFrameLease {
20    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        formatter.write_str("SubscriptionFrameLease")
22    }
23}
24
25/// One frame emitted toward the wire.
26#[derive(Debug, Clone)]
27pub enum PortalFrame {
28    /// Rows produced in a cycle.
29    Batch {
30        /// Arrow rows in the shared-log entry.
31        batch: RecordBatch,
32        /// Physical sequence within this in-memory object incarnation; not a durable resume token.
33        sequence: u64,
34        /// Internal process-memory ownership token.
35        #[doc(hidden)]
36        lease: SubscriptionFrameLease,
37    },
38    /// Progress frontier for a durably committed checkpoint.
39    Barrier {
40        /// Physical sequence of this progress entry within the current object incarnation.
41        sequence: u64,
42        /// Engine checkpoint epoch.
43        epoch: u64,
44        /// Engine checkpoint id.
45        checkpoint_id: u64,
46        /// Every shared-log entry with sequence below this value is covered by the cut.
47        through_sequence: u64,
48    },
49    /// Consumer fell behind by exactly `skipped` shared-log entries. This is
50    /// terminal because continuing would hide rows or checkpoint markers.
51    Lagged(u64),
52    /// The subscription cannot continue without returning an invalid result.
53    Error {
54        /// Human-readable failure detail.
55        message: String,
56    },
57}
58
59/// One `SUBSCRIBE` consumer.
60#[derive(Debug)]
61pub struct SubscriptionPortal {
62    name: String,
63    schema: SchemaRef,
64    reader: Option<SubscriptionReader>,
65    closed: bool,
66    filter: Option<Arc<dyn PhysicalExpr>>,
67}
68
69impl SubscriptionPortal {
70    pub(crate) fn open(
71        name: impl Into<String>,
72        schema: SchemaRef,
73        reader: SubscriptionReader,
74    ) -> Self {
75        Self::new_inner(name, schema, reader, None)
76    }
77
78    pub(crate) fn open_with_filter(
79        name: impl Into<String>,
80        schema: SchemaRef,
81        reader: SubscriptionReader,
82        filter: Arc<dyn PhysicalExpr>,
83    ) -> Self {
84        Self::new_inner(name, schema, reader, Some(filter))
85    }
86
87    fn new_inner(
88        name: impl Into<String>,
89        schema: SchemaRef,
90        reader: SubscriptionReader,
91        filter: Option<Arc<dyn PhysicalExpr>>,
92    ) -> Self {
93        Self {
94            name: name.into(),
95            schema,
96            reader: Some(reader),
97            closed: false,
98            filter,
99        }
100    }
101
102    /// Schema of the subscribed object.
103    #[must_use]
104    pub fn schema(&self) -> SchemaRef {
105        Arc::clone(&self.schema)
106    }
107
108    /// Next frame, or `None` after a terminal frame or explicit close.
109    pub async fn next_frame(&mut self) -> Option<PortalFrame> {
110        if self.closed {
111            return None;
112        }
113
114        loop {
115            let read = self.reader.as_mut()?.next().await;
116            if let Some(frame) = self.process_read(read) {
117                return Some(frame);
118            }
119        }
120    }
121
122    /// Return the next immediately available frame without waiting.
123    pub fn try_next_frame(&mut self) -> Option<PortalFrame> {
124        if self.closed {
125            return None;
126        }
127
128        loop {
129            let read = self.reader.as_mut()?.next().now_or_never()?;
130            if let Some(frame) = self.process_read(read) {
131                return Some(frame);
132            }
133        }
134    }
135
136    fn process_read(&mut self, read: SubscriptionRead) -> Option<PortalFrame> {
137        let frame = match read {
138            SubscriptionRead::Update { sequence, update } => translate(sequence, update),
139            SubscriptionRead::Lagged(skipped) => {
140                tracing::warn!(
141                    subscription = %self.name,
142                    skipped,
143                    "subscription cursor was evicted; closing"
144                );
145                self.close();
146                return Some(PortalFrame::Lagged(skipped));
147            }
148            SubscriptionRead::Terminal(message) => {
149                tracing::warn!(
150                    subscription = %self.name,
151                    %message,
152                    "subscription log terminated; closing"
153                );
154                self.close();
155                return Some(PortalFrame::Error { message });
156            }
157        };
158
159        let PortalFrame::Batch {
160            batch,
161            sequence,
162            lease,
163        } = frame
164        else {
165            if matches!(&frame, PortalFrame::Error { .. }) {
166                self.close();
167            }
168            return Some(frame);
169        };
170        let Some(filter) = self.filter.as_ref() else {
171            return Some(PortalFrame::Batch {
172                batch,
173                sequence,
174                lease,
175            });
176        };
177        match crate::filter_compile::apply(&batch, filter.as_ref()) {
178            Ok(Some(filtered)) => Some(PortalFrame::Batch {
179                batch: filtered,
180                sequence,
181                lease,
182            }),
183            Ok(None) => None,
184            Err(error) => {
185                tracing::warn!(
186                    subscription = %self.name,
187                    %error,
188                    "subscription filter failed; closing"
189                );
190                let message = error.to_string();
191                self.close();
192                Some(PortalFrame::Error { message })
193            }
194        }
195    }
196
197    /// Stop reading and release the subscriber registration. Idempotent.
198    pub fn close(&mut self) {
199        self.closed = true;
200        self.reader = None;
201    }
202
203    /// True after `close()` has been called.
204    #[must_use]
205    pub fn is_closed(&self) -> bool {
206        self.closed
207    }
208}
209
210impl Drop for SubscriptionPortal {
211    fn drop(&mut self) {
212        self.close();
213    }
214}
215
216fn translate(sequence: u64, update: ChargedUpdate) -> PortalFrame {
217    match update.as_ref() {
218        MvUpdate::Batch(batch) => {
219            let batch = batch.clone();
220            PortalFrame::Batch {
221                batch,
222                sequence,
223                lease: SubscriptionFrameLease { _owner: update },
224            }
225        }
226        MvUpdate::Barrier {
227            epoch,
228            checkpoint_id,
229            through_sequence,
230        } => PortalFrame::Barrier {
231            sequence,
232            epoch: *epoch,
233            checkpoint_id: *checkpoint_id,
234            through_sequence: *through_sequence,
235        },
236        MvUpdate::Error(message) => PortalFrame::Error {
237            message: message.clone(),
238        },
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use std::sync::Arc as StdArc;
245    use std::time::Duration;
246
247    use arrow_array::{Int64Array, RecordBatch};
248    use arrow_schema::{DataType, Field, Schema};
249
250    use super::super::registry::{
251        approx_size, MvUpdate, SubscribeStart, SubscriptionRegistry, MAX_LIVE_BATCH_BYTES,
252    };
253    use super::*;
254
255    fn schema() -> SchemaRef {
256        StdArc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]))
257    }
258
259    fn batch(ids: Vec<i64>) -> RecordBatch {
260        RecordBatch::try_new(schema(), vec![StdArc::new(Int64Array::from(ids))]).unwrap()
261    }
262
263    fn open(registry: &SubscriptionRegistry, start: SubscribeStart) -> SubscriptionPortal {
264        let reader = registry.subscribe("mv", start).unwrap();
265        SubscriptionPortal::open("mv", schema(), reader)
266    }
267
268    #[tokio::test]
269    async fn portal_forwards_batch_and_barrier() {
270        let registry = SubscriptionRegistry::new();
271        let mut portal = open(&registry, SubscribeStart::Tail);
272        registry.send_batch("mv", batch(vec![1, 2])).unwrap();
273        registry.broadcast_barrier(7, 7);
274
275        assert!(matches!(
276            portal.next_frame().await,
277            Some(PortalFrame::Batch { batch, sequence: 0, .. }) if batch.num_rows() == 2
278        ));
279        assert!(matches!(
280            portal.next_frame().await,
281            Some(PortalFrame::Barrier {
282                sequence: 1,
283                epoch: 7,
284                checkpoint_id: 7,
285                through_sequence: 1,
286            })
287        ));
288    }
289
290    #[test]
291    fn portal_try_next_is_non_blocking_and_ordered() {
292        let registry = SubscriptionRegistry::new();
293        let mut portal = open(&registry, SubscribeStart::Tail);
294        assert!(portal.try_next_frame().is_none());
295
296        registry.send_batch("mv", batch(vec![1])).unwrap();
297        registry.broadcast_barrier(3, 3);
298        assert!(matches!(
299            portal.try_next_frame(),
300            Some(PortalFrame::Batch { batch, sequence: 0, .. }) if batch.num_rows() == 1
301        ));
302        assert!(matches!(
303            portal.try_next_frame(),
304            Some(PortalFrame::Barrier {
305                sequence: 1,
306                epoch: 3,
307                checkpoint_id: 3,
308                through_sequence: 1,
309            })
310        ));
311        assert!(portal.try_next_frame().is_none());
312    }
313
314    #[tokio::test]
315    async fn portal_reports_object_drop_then_closes() {
316        let registry = SubscriptionRegistry::new();
317        let mut portal = open(&registry, SubscribeStart::Tail);
318        assert!(registry.drop_name("mv"));
319
320        let frame = tokio::time::timeout(Duration::from_millis(500), portal.next_frame())
321            .await
322            .unwrap();
323        assert!(matches!(
324            frame,
325            Some(PortalFrame::Error { message }) if message == "object dropped"
326        ));
327        assert!(portal.next_frame().await.is_none());
328    }
329
330    #[tokio::test]
331    async fn portal_emits_exact_lag_as_final_frame() {
332        let registry = SubscriptionRegistry::new();
333        let mut portal = open(&registry, SubscribeStart::Tail);
334        let rows = (MAX_LIVE_BATCH_BYTES / 2) / std::mem::size_of::<i64>();
335        for value in 0..6_i64 {
336            registry.send_batch("mv", batch(vec![value; rows])).unwrap();
337        }
338
339        assert!(matches!(
340            portal.next_frame().await,
341            Some(PortalFrame::Lagged(skipped)) if skipped > 0
342        ));
343        assert!(portal.next_frame().await.is_none());
344    }
345
346    #[tokio::test]
347    async fn portal_reads_shared_as_of_suffix_before_new_live_entries() {
348        let registry = SubscriptionRegistry::new();
349        registry.configure("mv", 1 << 20);
350        registry.broadcast_barrier(1, 1);
351        registry.send_batch("mv", batch(vec![10])).unwrap();
352        registry.broadcast_barrier(2, 2);
353        registry.send_batch("mv", batch(vec![20])).unwrap();
354        let mut portal = open(&registry, SubscribeStart::AsOfEpoch(1));
355        registry.send_batch("mv", batch(vec![30])).unwrap();
356
357        let mut rows = Vec::new();
358        let mut barriers = Vec::new();
359        for _ in 0..4 {
360            match portal.next_frame().await.unwrap() {
361                PortalFrame::Batch { batch, .. } => rows.push(
362                    batch
363                        .column(0)
364                        .as_any()
365                        .downcast_ref::<Int64Array>()
366                        .unwrap()
367                        .value(0),
368                ),
369                PortalFrame::Barrier { epoch, .. } => barriers.push(epoch),
370                PortalFrame::Lagged(skipped) => panic!("unexpected lag of {skipped}"),
371                PortalFrame::Error { message } => panic!("unexpected error: {message}"),
372            }
373        }
374        assert_eq!(rows, vec![10, 20, 30]);
375        assert_eq!(barriers, vec![2]);
376    }
377
378    #[test]
379    fn close_releases_registration_once() {
380        let registry = SubscriptionRegistry::new();
381        let mut portal = open(&registry, SubscribeStart::Tail);
382        assert_eq!(registry.subscriber_count("mv"), 1);
383        portal.close();
384        portal.close();
385        assert_eq!(registry.subscriber_count("mv"), 0);
386    }
387
388    #[tokio::test]
389    async fn held_batch_frame_blocks_process_budget_reuse() {
390        let sample = batch(vec![1, 2, 3]);
391        let entry_bytes = approx_size(&MvUpdate::Batch(sample.clone()));
392        let registry = SubscriptionRegistry::with_storage_budget(entry_bytes);
393        let mut portal = open(&registry, SubscribeStart::Tail);
394        registry.send_batch("mv", sample.clone()).unwrap();
395
396        let frame = portal.next_frame().await.unwrap();
397        assert!(matches!(&frame, PortalFrame::Batch { .. }));
398        assert_eq!(registry.charged_bytes(), entry_bytes);
399
400        registry.configure("contender", entry_bytes.saturating_mul(4));
401        assert!(registry.send_batch("contender", sample.clone()).is_err());
402        assert_eq!(registry.charged_bytes(), entry_bytes);
403
404        drop(frame);
405        assert_eq!(registry.charged_bytes(), 0);
406        registry.configure("replacement", entry_bytes.saturating_mul(4));
407        registry.send_batch("replacement", sample).unwrap();
408        assert_eq!(registry.charged_bytes(), entry_bytes);
409    }
410
411    #[tokio::test]
412    async fn two_portals_share_one_charge_until_both_frames_drop() {
413        let sample = batch(vec![1, 2, 3]);
414        let entry_bytes = approx_size(&MvUpdate::Batch(sample.clone()));
415        let registry = SubscriptionRegistry::with_storage_budget(entry_bytes);
416        let mut first = open(&registry, SubscribeStart::Tail);
417        let mut second = open(&registry, SubscribeStart::Tail);
418        registry.send_batch("mv", sample).unwrap();
419
420        let first_frame = first.next_frame().await.unwrap();
421        let second_frame = second.next_frame().await.unwrap();
422        assert_eq!(registry.charged_bytes(), entry_bytes);
423
424        drop(first_frame);
425        assert_eq!(registry.charged_bytes(), entry_bytes);
426        drop(second_frame);
427        assert_eq!(registry.charged_bytes(), 0);
428    }
429
430    #[tokio::test]
431    async fn filtered_batch_keeps_the_original_entry_charge() {
432        let sample = batch(vec![-1, 2]);
433        let entry_bytes = approx_size(&MvUpdate::Batch(sample.clone()));
434        let registry = SubscriptionRegistry::with_storage_budget(entry_bytes);
435        let reader = registry.subscribe("mv", SubscribeStart::Tail).unwrap();
436        let context = datafusion::prelude::SessionContext::new();
437        let filter = crate::filter_compile::compile(&context, "id > 0", &schema())
438            .await
439            .unwrap();
440        let mut portal = SubscriptionPortal::open_with_filter("mv", schema(), reader, filter);
441        registry.send_batch("mv", sample).unwrap();
442
443        let frame = portal.next_frame().await.unwrap();
444        assert!(matches!(
445            &frame,
446            PortalFrame::Batch { batch, .. } if batch.num_rows() == 1
447        ));
448        assert_eq!(registry.charged_bytes(), entry_bytes);
449
450        drop(frame);
451        assert_eq!(registry.charged_bytes(), 0);
452    }
453
454    #[tokio::test]
455    async fn object_drop_does_not_release_a_held_frame_charge() {
456        let sample = batch(vec![1]);
457        let entry_bytes = approx_size(&MvUpdate::Batch(sample.clone()));
458        let registry = SubscriptionRegistry::with_storage_budget(entry_bytes);
459        let mut portal = open(&registry, SubscribeStart::Tail);
460        registry.send_batch("mv", sample).unwrap();
461        let frame = portal.next_frame().await.unwrap();
462
463        assert!(registry.drop_name("mv"));
464        assert_eq!(registry.charged_bytes(), entry_bytes);
465
466        drop(frame);
467        assert_eq!(registry.charged_bytes(), 0);
468    }
469
470    #[tokio::test]
471    async fn filter_evaluation_failure_is_terminal_and_visible() {
472        let registry = SubscriptionRegistry::new();
473        let reader = registry.subscribe("mv", SubscribeStart::Tail).unwrap();
474        let context = datafusion::prelude::SessionContext::new();
475        let filter = crate::filter_compile::compile(&context, "id > 0", &schema())
476            .await
477            .unwrap();
478        let mut portal = SubscriptionPortal::open_with_filter("mv", schema(), reader, filter);
479
480        let incompatible_schema =
481            StdArc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)]));
482        let incompatible = RecordBatch::try_new(
483            incompatible_schema,
484            vec![StdArc::new(arrow_array::StringArray::from(vec!["bad"]))],
485        )
486        .unwrap();
487        registry.send_batch("mv", incompatible).unwrap();
488
489        assert!(matches!(
490            portal.next_frame().await,
491            Some(PortalFrame::Error { message }) if message.contains("filter")
492        ));
493        assert!(portal.next_frame().await.is_none());
494        assert_eq!(registry.subscriber_count("mv"), 0);
495    }
496}
laminar_db/subscription/portal.rs

laminar_db/subscription/
portal.rs
