harmos_runtime/streams/
batches.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::Any;
18use std::collections::VecDeque;
19use std::future::Future;
20use std::marker::PhantomData;
21use std::pin::Pin;
22use std::sync::{Arc, Mutex};
23use std::task::{Context, Poll};
24// Third Party Imports
25use harmos_domain::Aggregate;
26use harmos_storage::signal::Summary as Running;
27use tokio::sync::watch;
28// Local Imports
29use super::demand::Holding;
30use crate::aggregate;
31use crate::streams::BatchDefinition;
32
33/* === Constants === */
34
35/* === Macros === */
36
37/* === Type Aliases === */
38
39type Next = Pin<Box<dyn Future<Output = (Cursor, Option<BatchItem>)> + Send>>;
40
41/* === Structs === */
42
43// One measured field of a batch. `statistics` is the field's running
44// summary, absent when no finite value arrived or a statistic overflowed.
45#[derive(Clone, Debug, PartialEq)]
46pub struct Measurement {
47    pub field: &'static str,
48    pub statistics: Option<Running>,
49    pub invalid: u64,
50    pub overflow: bool,
51}
52
53#[derive(Clone, Debug, PartialEq)]
54pub struct Summary {
55    pub count: u64,
56    pub missed: u64,
57    pub partial: bool,
58    pub fields: Vec<Measurement>,
59}
60
61pub struct Batches<S> {
62    cursor: Option<Cursor>,
63    next: Option<Next>,
64    marker: PhantomData<fn() -> S>,
65    _held: Holding,
66}
67
68struct Cursor {
69    feed: Arc<Feed>,
70    bell: watch::Receiver<()>,
71    next: u64,
72}
73
74pub(crate) struct Feed {
75    ring: Mutex<Retained>,
76    bell: watch::Sender<()>,
77}
78
79struct Retained {
80    rows: VecDeque<Arc<Summary>>,
81    next: u64,
82    closed: bool,
83}
84
85pub(crate) struct State {
86    definition: BatchDefinition,
87    count: u64,
88    missed: u64,
89    fields: Vec<Accumulator>,
90    pub(crate) feed: Arc<Feed>,
91    closed: bool,
92}
93
94#[derive(Clone, Default)]
95struct Accumulator {
96    running: Running,
97    invalid: u64,
98}
99
100/* === Enums === */
101
102#[derive(Clone, Debug, PartialEq)]
103pub enum BatchItem {
104    Summary(Arc<Summary>),
105    Gap { missed: u64 },
106}
107
108/* === Traits === */
109
110/* === Implementations === */
111
112impl Summary {
113    pub fn complete(&self) -> bool {
114        !self.partial && self.missed == 0
115    }
116}
117
118impl<S> Batches<S> {
119    pub async fn next(&mut self) -> Option<BatchItem> {
120        std::future::poll_fn(|cx| futures_core::Stream::poll_next(Pin::new(&mut *self), cx)).await
121    }
122}
123
124impl Cursor {
125    async fn advance(mut self) -> (Self, Option<BatchItem>) {
126        loop {
127            let (item, closed) = {
128                let ring = self.feed.ring.lock().expect("summary ring is healthy");
129                let floor = ring.next - ring.rows.len() as u64;
130                let item = if self.next < floor {
131                    let missed = floor - self.next;
132                    self.next = floor;
133                    Some(BatchItem::Gap { missed })
134                } else if self.next < ring.next {
135                    let row = Arc::clone(&ring.rows[(self.next - floor) as usize]);
136                    self.next += 1;
137                    Some(BatchItem::Summary(row))
138                } else {
139                    None
140                };
141                (item, ring.closed)
142            };
143            if item.is_some() || closed {
144                return (self, item);
145            }
146            if self.bell.changed().await.is_err() {
147                return (self, None);
148            }
149        }
150    }
151}
152
153impl<S> futures_core::Stream for Batches<S> {
154    type Item = BatchItem;
155    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
156        let this = self.get_mut();
157        if this.next.is_none() {
158            let Some(cursor) = this.cursor.take() else {
159                return Poll::Ready(None);
160            };
161            this.next = Some(Box::pin(cursor.advance()));
162        }
163        match this.next.as_mut().expect("advance installed").as_mut().poll(cx) {
164            Poll::Pending => Poll::Pending,
165            Poll::Ready((cursor, item)) => {
166                this.next = None;
167                if item.is_some() {
168                    this.cursor = Some(cursor);
169                }
170                Poll::Ready(item)
171            }
172        }
173    }
174}
175
176impl Feed {
177    pub(crate) fn following<S>(self: &Arc<Self>, held: Holding) -> Batches<S> {
178        let bell = self.bell.subscribe();
179        let next = self.ring.lock().expect("summary ring is healthy").next;
180        Batches { cursor: Some(Cursor { feed: Arc::clone(self), bell, next }), next: None, marker: PhantomData, _held: held }
181    }
182}
183
184impl State {
185    pub(crate) fn new(definition: BatchDefinition) -> Self {
186        let (bell, _) = watch::channel(());
187        Self {
188            definition,
189            count: 0,
190            missed: 0,
191            fields: vec![Accumulator::default(); definition.fields().len()],
192            feed: Arc::new(Feed { ring: Mutex::new(Retained { rows: VecDeque::with_capacity(64), next: 0, closed: false }), bell }),
193            closed: false,
194        }
195    }
196
197    pub(crate) fn publish(&mut self, row: &dyn Any) {
198        if self.closed {
199            return;
200        }
201        let mut fields = self.fields.iter_mut();
202        self.definition.measure(row, &mut |value| fields.next().expect("one accumulator per declared measure").push(value));
203        self.count += 1;
204        if self.count == self.definition.samples() {
205            self.flush(false);
206        }
207    }
208
209    pub(crate) fn miss(&mut self, count: u64) {
210        if !self.closed {
211            self.missed = self.missed.saturating_add(count);
212        }
213    }
214
215    fn flush(&mut self, partial: bool) {
216        let summary = Summary {
217            count: self.count,
218            missed: self.missed,
219            partial,
220            fields: self
221                .fields
222                .iter_mut()
223                .zip(self.definition.fields())
224                .map(|(accumulator, field)| {
225                    let result = accumulator.finish(field);
226                    *accumulator = Accumulator::default();
227                    result
228                })
229                .collect(),
230        };
231        self.count = 0;
232        self.missed = 0;
233        let mut ring = self.feed.ring.lock().expect("summary ring is healthy");
234        if ring.rows.len() == 64 {
235            ring.rows.pop_front();
236        }
237        ring.rows.push_back(Arc::new(summary));
238        ring.next += 1;
239        drop(ring);
240        self.feed.bell.send_replace(());
241    }
242
243    pub(crate) fn close(&mut self) {
244        if self.closed {
245            return;
246        }
247        self.closed = true;
248        if self.count > 0 || self.missed > 0 {
249            self.flush(true);
250        }
251        self.feed.ring.lock().expect("summary ring is healthy").closed = true;
252        self.feed.bell.send_replace(());
253    }
254}
255
256impl Measurement {
257    // Mergeable members only, as `Running` answers them; a batch keeps no
258    // values, so median and percentiles answer `None`.
259    pub fn aggregate(&self, kind: Aggregate) -> Option<f64> {
260        aggregate::of_summary(self.statistics.as_ref()?, kind)
261    }
262}
263
264impl Accumulator {
265    fn push(&mut self, value: f64) {
266        if value.is_finite() {
267            self.running.push(value);
268        } else {
269            self.invalid += 1;
270        }
271    }
272
273    fn finish(&self, field: &'static str) -> Measurement {
274        let running = self.running;
275        let overflow = !running.is_empty()
276            && [Some(running.sum()), running.mean(), running.variance(), running.rms()]
277                .into_iter()
278                .any(|statistic| !statistic.is_some_and(f64::is_finite));
279        let statistics = (!running.is_empty() && !overflow).then_some(running);
280        Measurement { field, statistics, invalid: self.invalid, overflow }
281    }
282}
283
284/* === Public Functions === */
285
286/* === Private Functions === */
287
288/* === Tests === */
harmos_runtime/streams/batches.rs

harmos_runtime/streams/
batches.rs
