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

harmos_runtime/streams/
manager.rs
