harmos_storage/
follow.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10// Live Storage Subscriptions
11
12/* === Modules === */
13
14/* === Re-Exports === */
15
16/* === Imports === */
17
18// Standard Library Imports
19use std::num::NonZeroUsize;
20use std::sync::Arc;
21use std::sync::atomic::{AtomicU64, Ordering};
22
23// Third Party Imports
24
25// Local Imports
26use super::backend::{ReadLimits, ScanPage};
27use super::capacity::{Attempt, Held, Pool, Refusal, Reservation};
28use super::config::SCAN_ROWS;
29use super::manager::{Domain, ENTRY_OVERHEAD, Inner, State, encoded_size};
30pub use super::selection::Entry;
31use super::selection::Queued;
32use super::snapshot::Snapshot;
33use super::{StorageCursor, StorageManager};
34use crate::{Error, Result};
35
36/* === Constants === */
37
38static NEXT_FOLLOWER: AtomicU64 = AtomicU64::new(1);
39
40/* === Macros === */
41
42/* === Type Aliases === */
43
44/* === Structs === */
45
46pub struct Follower<T: Entry> {
47    context: crate::execution::Context,
48    inner: Arc<Inner>,
49    id: u64,
50    held: Held,
51    selection: T::Selection,
52    history: Option<Snapshot>,
53    after: u64,
54    saved: StorageCursor,
55    accepted: StorageCursor,
56}
57
58pub struct FollowPage<T> {
59    entries: Vec<T>,
60    _charge: crate::arrow::RetainedMemory,
61}
62
63pub(super) struct Following {
64    pub(super) after: u64,
65    failure: Option<Failure>,
66}
67
68/* === Enums === */
69
70#[derive(Clone, Copy)]
71enum Failure {
72    Gap { requested: u64, available: u64 },
73    Closed,
74}
75
76/* === Traits === */
77
78/* === Implementations === */
79
80impl<C> StorageManager<C> {
81    #[tracing::instrument(name = "storage.follow.follow", level = "debug", skip_all)]
82    pub async fn follow<T: Entry>(&self, after: Option<&StorageCursor>, selection: T::Selection) -> Result<Follower<T>> {
83        if let Some(after) = after {
84            self.inner.primary.history().check(&after.history)?;
85        }
86        let requested = after.map_or(0, |after| after.sequence);
87        let after = requested.max(T::start(&selection));
88        let id = NEXT_FOLLOWER.fetch_add(1, Ordering::Relaxed);
89
90        // Register Queue Retention From The Saved Frontier Atomically.
91        let (saved, accepted) = {
92            let mut state = self.inner.state.lock().unwrap();
93            super::capacity::admitting(&state, Pool::Read)?;
94            let saved = state.status.saved.clone();
95            let accepted = state.status.accepted.clone();
96            if requested > accepted.sequence {
97                return Err(Error::InvalidInput { field: "storage follow", reason: "start exceeds accepted boundary" });
98            }
99            state.followers.insert(id, Following { after: saved.sequence.max(after), failure: None });
100            (saved, accepted)
101        };
102        let mut follower = Follower {
103            context: crate::execution::Context::current(),
104            inner: Arc::clone(&self.inner),
105            id,
106            held: Held::default(),
107            selection,
108            history: None,
109            after,
110            saved,
111            accepted,
112        };
113        if after < follower.saved.sequence {
114            follower.history = Some(self.snapshot(&follower.saved, &follower.held).await?);
115        }
116        Ok(follower)
117    }
118}
119
120impl<C> StorageManager<C> {
121    #[tracing::instrument(name = "storage.follow.page", level = "debug", skip_all)]
122    pub async fn page<T: Entry>(&self, selection: &T::Selection, after: u64, limits: ReadLimits) -> Result<FollowPage<T>> {
123        let held = Held::default();
124        let rows = limits.rows.min(SCAN_ROWS);
125        let bytes = limits.bytes.min(self.inner.settings.scan_bytes());
126        let charge = rows.get().saturating_mul(ENTRY_OVERHEAD).saturating_add(bytes.get());
127        let reservation = Reservation::wait(&self.inner, Pool::Read, &held, charge, None).await.map_err(Refusal::into_error)?.retain();
128        let snapshot = self.snapshot(&self.cursor_saved(), &held).await?;
129        let (entries, encoded) = snapshot.scan::<T>(selection, ScanPage { after, rows, bytes }, reservation.clone()).await?;
130        reservation.shrink_to(encoded.saturating_add(entries.len().saturating_mul(ENTRY_OVERHEAD)))?;
131        Ok(FollowPage { entries, _charge: reservation })
132    }
133}
134
135impl<T: Entry> Follower<T> {
136    pub fn saved_boundary(&self) -> &StorageCursor {
137        &self.saved
138    }
139
140    pub fn accepted_boundary(&self) -> &StorageCursor {
141        &self.accepted
142    }
143
144    pub async fn next(&mut self) -> Result<FollowPage<T>> {
145        let context = self.context.clone();
146        let span = context.sync(|| tracing::trace_span!("storage.reader.page"));
147        let future = async {
148            loop {
149                if let Some(page) = self.next_history().await? {
150                    return Ok(page);
151                }
152                if let Some(page) = self.next_live().await? {
153                    return Ok(page);
154                }
155            }
156        };
157        context.scope(tracing::Instrument::instrument(future, span)).await
158    }
159
160    // Saved History Through The Registration Frontier, From The Pinned Image.
161    #[tracing::instrument(name = "storage.follow.next_history", level = "debug", skip_all)]
162    async fn next_history(&mut self) -> Result<Option<FollowPage<T>>> {
163        let Some(snapshot) = &self.history else {
164            return Ok(None);
165        };
166        let bytes = self.inner.settings.scan_bytes();
167        let reservation =
168            Reservation::wait(&self.inner, Pool::Read, &self.held, page_charge(bytes), None).await.map_err(Refusal::into_error)?.retain();
169        let page = ScanPage { after: self.after, rows: SCAN_ROWS, bytes };
170        let (entries, encoded) = snapshot.scan::<T>(&self.selection, page, reservation.clone()).await?;
171        let Some(last) = entries.last() else {
172            self.history = None;
173            self.after = self.after.max(self.saved.sequence);
174            return Ok(None);
175        };
176        self.after = last.position().sequence;
177        reservation.shrink_to(encoded.saturating_add(entries.len().saturating_mul(ENTRY_OVERHEAD)))?;
178        Ok(Some(FollowPage { entries, _charge: reservation }))
179    }
180
181    // Accepted Entries From The Queue; Waits Without Allocating While None Match.
182    #[tracing::instrument(name = "storage.follow.next_live", level = "debug", skip_all)]
183    async fn next_live(&mut self) -> Result<Option<FollowPage<T>>> {
184        let id = self.id;
185        let selection = &self.selection;
186        self.inner
187            .wait(None, |state| match find::<T>(state, id, selection) {
188                Ok(true) => Attempt::Ready(()),
189                Ok(false) => Attempt::Wait,
190                Err(error) => Attempt::Refuse(Refusal::Failed(error)),
191            })
192            .await
193            .map_err(Refusal::into_error)?;
194
195        // Charge The Page Before Copying Entries Out Of The Queue.
196        let bytes = self.inner.settings.scan_bytes();
197        let mut reservation =
198            Reservation::wait(&self.inner, Pool::Read, &self.held, page_charge(bytes), None).await.map_err(Refusal::into_error)?;
199        let (entries, encoded) = {
200            let mut state = self.inner.state.lock().unwrap();
201            copy::<T>(&mut state, self.id, &self.selection, bytes)?
202        };
203        let Some(last) = entries.last() else {
204            return Ok(None);
205        };
206        self.after = last.position().sequence;
207        reservation.shrink_to(encoded.saturating_add(entries.len().saturating_mul(ENTRY_OVERHEAD)));
208        Ok(Some(FollowPage { entries, _charge: reservation.retain() }))
209    }
210}
211
212impl<T: Entry> Drop for Follower<T> {
213    fn drop(&mut self) {
214        let mut state = self.inner.state.lock().unwrap();
215        state.followers.remove(&self.id);
216        state.retire_data();
217        state.collect();
218        drop(state);
219        self.inner.notify();
220    }
221}
222
223impl<T> FollowPage<T> {
224    pub fn entries(&self) -> &[T] {
225        &self.entries
226    }
227}
228
229impl Following {
230    fn retains(&self) -> bool {
231        self.failure.is_none()
232    }
233
234    fn check(&self) -> Result<()> {
235        match self.failure {
236            Some(Failure::Gap { requested, available }) => Err(Error::Gap { requested, available }),
237            Some(Failure::Closed) => Err(Error::Closing { domain: "storage" }),
238            None => Ok(()),
239        }
240    }
241}
242
243impl State {
244    pub(super) fn close_followers(&mut self) {
245        for following in self.followers.values_mut() {
246            following.failure.get_or_insert(Failure::Closed);
247        }
248        self.retire_data();
249    }
250
251    pub(super) fn retire_data(&mut self) {
252        let through =
253            self.followers.values().filter(|following| following.retains()).map(|following| following.after).min().unwrap_or(u64::MAX);
254        self.queue.retire_data(self.data_flushed, through);
255    }
256
257    pub(super) fn lag_followers(&mut self) -> bool {
258        let Some(through) = self
259            .queue
260            .iter()
261            .filter(|entry| entry.operation.domain() == Domain::Data && entry.index <= self.data_flushed)
262            .map(|entry| entry.operation.sequence())
263            .max()
264        else {
265            return false;
266        };
267        let mut changed = false;
268        for following in self.followers.values_mut() {
269            if following.retains() && following.after < through {
270                following.failure =
271                    Some(Failure::Gap { requested: following.after.saturating_add(1), available: through.saturating_add(1) });
272                changed = true;
273            }
274        }
275        if changed {
276            self.retire_data();
277            self.collect();
278        }
279        changed
280    }
281}
282
283/* === Public Functions === */
284
285/* === Private Functions === */
286
287fn page_charge(bytes: NonZeroUsize) -> usize {
288    SCAN_ROWS.get().saturating_mul(ENTRY_OVERHEAD).saturating_add(bytes.get())
289}
290
291// Whether A Matching Queued Entry Follows The Subscription Cursor. Queue
292// Entries Without A Match Advance The Cursor So They Can Be Retired.
293fn find<T: Entry>(state: &mut State, id: u64, selection: &T::Selection) -> Result<bool> {
294    super::capacity::admitting(state, Pool::Read)?;
295    let following = state.followers.get(&id).ok_or(Error::Closing { domain: "storage follow" })?;
296    following.check()?;
297    let mut cursor = following.after;
298    let mut found = false;
299    for pending in state.queue.iter().filter(|entry| entry.operation.domain() == Domain::Data) {
300        if pending.operation.sequence() <= cursor {
301            continue;
302        }
303        let entries = T::pending(Queued::new(&pending.operation));
304        let start = entries.partition_point(|entry| entry.position().sequence <= cursor);
305        if entries[start..].iter().any(|entry| entry.matches(selection)) {
306            found = true;
307            break;
308        }
309        cursor = pending.operation.sequence();
310    }
311    let following = state.followers.get_mut(&id).expect("registered follower");
312    if cursor > following.after {
313        following.after = cursor;
314        state.retire_data();
315        state.collect();
316    }
317    Ok(found)
318}
319
320// Copy Matching Queued Entries After The Cursor Within The Page Bounds.
321fn copy<T: Entry>(state: &mut State, id: u64, selection: &T::Selection, bytes: NonZeroUsize) -> Result<(Vec<T>, usize)> {
322    let following = state.followers.get(&id).ok_or(Error::Closing { domain: "storage follow" })?;
323    following.check()?;
324    let mut cursor = following.after;
325    let mut entries = Vec::new();
326    let mut encoded = 0usize;
327    'queue: for pending in state.queue.iter().filter(|entry| entry.operation.domain() == Domain::Data) {
328        if pending.operation.sequence() <= cursor {
329            continue;
330        }
331        let values = T::pending(Queued::new(&pending.operation));
332        let start = values.partition_point(|entry| entry.position().sequence <= cursor);
333        for entry in &values[start..] {
334            if entries.len() >= SCAN_ROWS.get() {
335                break 'queue;
336            }
337            if entry.matches(selection) {
338                let size = encoded_size(entry)?;
339                if !entries.is_empty() && encoded.saturating_add(size) > bytes.get() {
340                    break 'queue;
341                }
342                encoded = encoded.saturating_add(size);
343                entries.push(entry.clone());
344            }
345            cursor = entry.position().sequence;
346        }
347        cursor = cursor.max(pending.operation.sequence());
348    }
349    let following = state.followers.get_mut(&id).expect("registered follower");
350    following.after = following.after.max(cursor);
351    state.retire_data();
352    state.collect();
353    Ok((entries, encoded))
354}
355
356/* === Tests === */
harmos_storage/follow.rs

harmos_storage/
follow.rs
