harmos_runtime/streams/
manager.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14/* === Imports === */
15
16// Standard Library Imports
17use std::any::TypeId;
18use std::fmt;
19use std::future::{Future, IntoFuture};
20use std::marker::PhantomData;
21use std::ops::{Bound, RangeBounds};
22use std::pin::Pin;
23use std::sync::Arc;
24use std::sync::atomic::{AtomicBool, Ordering};
25// Third Party Imports
26use harmos_domain::Aggregate;
27use tokio::sync::watch;
28// Local Imports
29use super::window::Window;
30use super::{
31    Aggregated, Axial, Batched, Batches, BeneathWindow, Bucket, Definition, Demand, Holding, Item, Prefix, Stream, Table, Tier, Tracked,
32    Wired, demand, window,
33};
34use crate::{Application, state::At, state::Position, storage::Batch, storage::Cursor};
35
36/* === Constants === */
37
38/* === Macros === */
39
40/* === Type Aliases === */
41
42/* === Structs === */
43
44pub struct StreamManager<A: Application> {
45    pub(super) schema_settings: Arc<crate::config::SchemaSettings>,
46    pub(super) windows: Arc<Table>,
47    pub(super) applied: watch::Receiver<Position>,
48    pub(super) tier: Tier,
49    pub(super) closed: Arc<AtomicBool>,
50    pub(super) marker: PhantomData<fn() -> A>,
51}
52
53// A live feed paces plugin publication: the window never evicts a row this
54// feed has not taken in. Asking for the next batch takes in the previous one.
55pub struct StreamFeed<A: Application> {
56    pub(super) streams: StreamManager<A>,
57    pub(super) changed: watch::Receiver<()>,
58    pub(super) cursor: Cursor,
59    pub(super) intake: super::ingest::Intake,
60}
61
62pub(super) struct Partition<S> {
63    pub(super) labels: Vec<String>,
64    pub(super) matches: Box<dyn Fn(&S) -> bool + Send>,
65}
66
67pub struct Query<S> {
68    pub(super) window: Arc<Window>,
69    pub(super) definition: Definition,
70    pub(super) tier: Tier,
71    pub(super) start: Bound<Position>,
72    pub(super) end: Bound<Position>,
73    pub(super) as_of: Option<Position>,
74    pub(super) partition: Option<Partition<S>>,
75    pub(super) retained_only: bool,
76    pub(super) marker: PhantomData<fn() -> S>,
77}
78
79pub struct Survey<S> {
80    pub(super) definition: Definition,
81    pub(super) tier: Tier,
82    pub(super) start: Bound<u64>,
83    pub(super) end: Bound<u64>,
84    pub(super) track: Vec<String>,
85    pub(super) marker: PhantomData<fn() -> S>,
86}
87
88pub struct Seal<S> {
89    pub(super) definition: Definition,
90    pub(super) tier: Tier,
91    pub(super) track: Vec<String>,
92    pub(super) marker: PhantomData<fn() -> S>,
93}
94
95/* === Enums === */
96
97/* === Traits === */
98
99/* === Implementations === */
100
101impl<A: Application> StreamFeed<A> {
102    pub async fn next(&mut self) -> Option<Batch<A>> {
103        self.intake.took(self.cursor.0);
104        loop {
105            self.changed.borrow_and_update();
106            let batch = self.streams.batch(self.cursor);
107            if batch.reaches() > self.cursor {
108                self.cursor = batch.reaches();
109                return Some(batch);
110            }
111            if self.streams.closed.load(Ordering::Acquire) {
112                return None;
113            }
114            if self.changed.changed().await.is_err() {
115                return None;
116            }
117        }
118    }
119}
120
121impl<A: Application> StreamManager<A> {
122    pub(crate) fn wired(
123        windows: Table, applied: watch::Receiver<Position>, tier: Tier, schema_settings: Arc<crate::config::SchemaSettings>,
124    ) -> Self {
125        Self {
126            schema_settings,
127            windows: Arc::new(windows),
128            applied,
129            tier,
130            closed: Arc::new(AtomicBool::new(false)),
131            marker: PhantomData,
132        }
133    }
134
135    pub fn follow(&self, resume: Cursor) -> StreamFeed<A> {
136        StreamFeed {
137            streams: self.clone(),
138            changed: self.delivery(),
139            cursor: resume,
140            intake: self.windows.ingest.engage(resume.0),
141        }
142    }
143
144    pub fn frontier(&self) -> Cursor {
145        self.windows.frontier()
146    }
147
148    pub fn publish<S: Stream<A>>(&self, row: S) {
149        let at = *self.applied.borrow();
150        self.stream::<S>().window.publish(at, Arc::new(row));
151    }
152
153    #[cfg(feature = "plugin-sidecar")]
154    pub(crate) fn miss<S: Stream<A>>(&self, count: u64) {
155        self.stream::<S>().window.miss(count);
156    }
157
158    pub fn query<S: Stream<A>>(&self, range: impl RangeBounds<Position>) -> Query<S> {
159        let wired = self.stream::<S>();
160
161        Query {
162            window: Arc::clone(&wired.window),
163            definition: wired.definition.clone(),
164            tier: self.tier.clone(),
165            start: range.start_bound().cloned(),
166            end: range.end_bound().cloned(),
167            as_of: None,
168            partition: None,
169            retained_only: false,
170            marker: PhantomData,
171        }
172    }
173
174    pub fn survey<S: Stream<A> + Axial>(&self, range: impl RangeBounds<u64>) -> Survey<S> {
175        Survey {
176            definition: self.stream::<S>().definition.clone(),
177            tier: self.tier.clone(),
178            start: range.start_bound().cloned(),
179            end: range.end_bound().cloned(),
180            track: Vec::new(),
181            marker: PhantomData,
182        }
183    }
184
185    pub fn seal<S: Stream<A>>(&self) -> Seal<S> {
186        Seal {
187            definition: self.stream::<S>().definition.clone(),
188            tier: self.tier.clone(),
189            track: Vec::new(),
190            marker: PhantomData,
191        }
192    }
193
194    pub fn watch<S: Stream<A>>(&self) -> impl futures_core::Stream<Item = Item<S>> + Send + Unpin + use<S, A> {
195        let window = Arc::clone(&self.stream::<S>().window);
196        let held = demand::Holding::engage(Arc::clone(&window));
197
198        window::following(window, Arc::clone(&self.closed), held)
199    }
200
201    pub fn batches<S: Stream<A> + Batched>(&self) -> Batches<S> {
202        let window = &self.stream::<S>().window;
203        window.batches::<S>(Holding::engage(Arc::clone(window)))
204    }
205
206    pub fn demand<S: Stream<A>>(&self) -> Demand {
207        Demand::of(&self.stream::<S>().window)
208    }
209
210    pub(crate) fn close(&self) {
211        for wired in self.windows.declared() {
212            wired.window.close_batches();
213        }
214        self.closed.store(true, Ordering::Release);
215        self.windows.wake();
216        self.windows.ingest.wake();
217    }
218
219    pub(crate) fn delivery(&self) -> watch::Receiver<()> {
220        self.windows.delivery()
221    }
222
223    pub(crate) fn batch(&self, resume: Cursor) -> Batch<A> {
224        let delivery = self.windows.recorded_after(resume);
225
226        Batch {
227            rows: delivery.rows,
228            incomplete: delivery.incomplete,
229            reaches: delivery.reaches,
230            application: PhantomData,
231        }
232    }
233
234    pub(crate) fn declared(&self) -> Vec<Arc<Wired>> {
235        self.windows.declared()
236    }
237
238    fn stream<S: Stream<A>>(&self) -> Arc<Wired> {
239        self.windows
240            .wired(S::ID, S::VERSION)
241            .filter(|wired| wired.definition.row() == TypeId::of::<S>() && wired.definition.version() == S::VERSION)
242            .unwrap_or_else(|| {
243                panic!(
244                    "the stream \"{}\" is not in the registered catalog; \
245                     register its definition with builder.register(catalog)",
246                    S::ID
247                )
248            })
249    }
250}
251
252impl<A: Application> Clone for StreamManager<A> {
253    fn clone(&self) -> Self {
254        Self {
255            schema_settings: Arc::clone(&self.schema_settings),
256            windows: Arc::clone(&self.windows),
257            applied: self.applied.clone(),
258            tier: self.tier.clone(),
259            closed: Arc::clone(&self.closed),
260            marker: PhantomData,
261        }
262    }
263}
264
265impl<A: Application> fmt::Debug for StreamManager<A> {
266    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
267        formatter.debug_struct("StreamManager").finish_non_exhaustive()
268    }
269}
270
271impl<S: Clone + Send + 'static> Query<S> {
272    pub fn retained(mut self) -> Self {
273        self.retained_only = true;
274        self
275    }
276
277    pub fn as_of(mut self, at: Position) -> Self {
278        self.as_of = Some(at);
279        self
280    }
281
282    async fn answer(self) -> Result<std::vec::IntoIter<At<S>>, BeneathWindow> {
283        let Self {
284            window,
285            definition,
286            tier,
287            start,
288            end,
289            as_of,
290            partition,
291            retained_only,
292            ..
293        } = self;
294        let retained = if retained_only { window.snapshot_scope(start, end, as_of, false) } else { window.snapshot(start, end, as_of) };
295        let mut rows = Vec::new();
296
297        if retained.beneath && !retained_only {
298            let served =
299                tier.rows(definition, narrowed(partition.as_ref()), start, beneath(beneath_of(end, as_of), retained.floor)).await?;
300
301            rows.extend(served.into_iter().map(|At { value, position }| (position, value)));
302        }
303        rows.extend(retained.rows);
304
305        Ok(rows
306            .into_iter()
307            .map(|(position, row)| At {
308                value: super::row_value::<S>(row.as_ref()).expect("a stream's tiers hold only its own rows").clone(),
309                position,
310            })
311            .filter(|At { value, .. }| partition.as_ref().is_none_or(|partition| (partition.matches)(value)))
312            .collect::<Vec<_>>()
313            .into_iter())
314    }
315}
316
317impl<S: Clone + Send + Tracked + 'static> Query<S> {
318    pub fn track<K: Prefix<S::Track>>(mut self, track: K) -> Self {
319        self.partition = Some(Partition {
320            labels: track.labels(),
321            matches: Box::new(move |row: &S| track.selects(&row.track())),
322        });
323        self
324    }
325}
326
327impl<S: Clone + Send + 'static> IntoFuture for Query<S> {
328    type Output = Result<std::vec::IntoIter<At<S>>, BeneathWindow>;
329    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;
330
331    fn into_future(self) -> Self::IntoFuture {
332        Box::pin(self.answer())
333    }
334}
335
336impl<S> Survey<S> {
337    pub async fn budget(self, points: usize) -> Result<Vec<Bucket>, BeneathWindow> {
338        self.tier.buckets(self.definition, self.track, self.start, self.end, points).await
339    }
340
341    pub async fn fold(self, kind: Aggregate) -> Result<Vec<Aggregated>, BeneathWindow> {
342        self.tier.aggregate(self.definition, self.track, self.start, self.end, checked(kind)).await
343    }
344}
345
346impl<S: Tracked> Survey<S> {
347    pub fn track<K: Prefix<S::Track>>(mut self, track: K) -> Self {
348        self.track = track.labels();
349        self
350    }
351}
352
353impl<S: Tracked> Seal<S> {
354    pub fn track<K: Prefix<S::Track>>(mut self, track: K) -> Self {
355        self.track = track.labels();
356        self
357    }
358}
359
360impl<S> IntoFuture for Seal<S> {
361    type Output = Position;
362    type IntoFuture = Pin<Box<dyn Future<Output = Position> + Send>>;
363
364    fn into_future(self) -> Self::IntoFuture {
365        Box::pin(async move { self.tier.seal(self.definition, self.track).await })
366    }
367}
368
369/* === Public Functions === */
370
371/* === Private Functions === */
372
373fn checked(kind: Aggregate) -> Aggregate {
374    assert!(kind.answerable(), "a percentile rank is a percentage, and {kind:?} names none");
375    kind
376}
377
378fn narrowed<S>(partition: Option<&Partition<S>>) -> Vec<String> {
379    partition.map(|partition| partition.labels.clone()).unwrap_or_default()
380}
381
382fn beneath(end: Bound<Position>, ceiling: Position) -> Bound<Position> {
383    match end {
384        Bound::Included(at) | Bound::Excluded(at) if at <= ceiling => end,
385        _ => Bound::Included(ceiling),
386    }
387}
388
389fn beneath_of(end: Bound<Position>, as_of: Option<Position>) -> Bound<Position> {
390    match as_of {
391        Some(cutoff) => beneath(end, cutoff),
392        None => end,
393    }
394}
395
396/* === Tests === */
harmos_runtime/streams/manager.rs

harmos_runtime/streams/
manager.rs
