harmos_runtime/state/
capture.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::collections::HashMap;
18use std::future::Future;
19use std::pin::Pin;
20use std::sync::{
21    Arc,
22    atomic::{AtomicBool, Ordering},
23};
24// Third Party Imports
25use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, watch};
26use tracing::Instrument;
27// Local Imports
28use crate::supervision::executor::InContext;
29use crate::{Application, state::Entry, state::Position, state::PreparationError, state::Receipt, state::Recorded, state::RequestKey};
30
31/* === Constants === */
32
33/* === Macros === */
34
35/* === Type Aliases === */
36
37/* === Structs === */
38
39#[derive(Clone, Copy, Debug)]
40pub struct CaptureLimits {
41    pub pending_bytes: u32,
42    pub max_group_bytes: u32,
43}
44
45pub(crate) struct CallbackDeclaration<A: Application> {
46    pub limits: CaptureLimits,
47    pub consumer: Box<dyn Capturing<A>>,
48}
49
50struct Group<A: Application> {
51    trace: crate::supervision::executor::SpanContext,
52    entries: Vec<Arc<Entry<Recorded<A>>>>,
53    frontier: Position,
54    _permit: OwnedSemaphorePermit,
55}
56
57pub(crate) struct Callback<A: Application> {
58    sender: mpsc::UnboundedSender<Group<A>>,
59    shared: Arc<Shared>,
60    limits: CaptureLimits,
61    worker: tokio::task::JoinHandle<()>,
62}
63
64struct Shared {
65    failed: Arc<AtomicBool>,
66    signal: watch::Sender<bool>,
67    capacity: Arc<Semaphore>,
68    progress: watch::Sender<Result<Position, CaptureError>>,
69}
70
71struct Failure {
72    shared: Arc<Shared>,
73    clean: bool,
74}
75
76/* === Enums === */
77
78pub(crate) enum Declaration<A: Application> {
79    Callback(CallbackDeclaration<A>),
80    Managed {
81        storage: crate::storage::StorageManager<crate::storage::ErasedCatalog>,
82        receipts: HashMap<RequestKey, Receipt>,
83        projection: Option<crate::storage::capture::Projection>,
84        snapshots: Option<crate::state::snapshot::Snapshots<A>>,
85    },
86}
87
88pub(crate) enum Capture<A: Application> {
89    Callback(Callback<A>),
90    Managed {
91        storage: crate::storage::StorageManager<crate::storage::ErasedCatalog>,
92        receipts: HashMap<RequestKey, Receipt>,
93        projection: Option<crate::storage::capture::Projection>,
94        snapshots: Option<crate::state::snapshot::Snapshots<A>>,
95        failed: Arc<AtomicBool>,
96        progress: watch::Sender<Result<Position, CaptureError>>,
97        signal: watch::Sender<bool>,
98    },
99}
100
101pub(crate) enum Permit {
102    Callback(OwnedSemaphorePermit),
103    Managed(Box<crate::storage::state::Permit>),
104}
105
106#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
107pub enum CaptureError {
108    #[error("required journal capture failed")]
109    Failed,
110
111    #[error("capture limits must be positive and group bytes must fit pending bytes")]
112    InvalidLimits,
113}
114
115/* === Traits === */
116
117pub trait JournalCapture<A: Application>: Send + 'static {
118    fn capture(
119        &mut self, entries: Vec<Arc<Entry<Recorded<A>>>>, frontier: Position,
120    ) -> impl Future<Output = Result<(), CaptureError>> + Send;
121}
122
123pub(crate) trait Capturing<A: Application>: Send {
124    fn capture(
125        &mut self, entries: Vec<Arc<Entry<Recorded<A>>>>, frontier: Position,
126    ) -> Pin<Box<dyn Future<Output = Result<(), CaptureError>> + Send + '_>>;
127}
128
129/* === Implementations === */
130
131impl<A: Application, C: JournalCapture<A>> Capturing<A> for C {
132    fn capture(
133        &mut self, entries: Vec<Arc<Entry<Recorded<A>>>>, frontier: Position,
134    ) -> Pin<Box<dyn Future<Output = Result<(), CaptureError>> + Send + '_>> {
135        Box::pin(JournalCapture::capture(self, entries, frontier))
136    }
137}
138
139impl Shared {
140    fn fail(&self) {
141        self.failed.store(true, Ordering::Release);
142        self.capacity.close();
143        let _ = self.progress.send_replace(Err(CaptureError::Failed));
144        self.signal.send_replace(true);
145    }
146}
147
148impl Drop for Failure {
149    fn drop(&mut self) {
150        if !self.clean {
151            self.shared.fail();
152        }
153    }
154}
155
156impl<A: Application> CallbackDeclaration<A> {
157    pub fn start(self, failed: Arc<AtomicBool>, origin: Position) -> Callback<A> {
158        let (sender, mut receiver) = mpsc::unbounded_channel::<Group<A>>();
159        let (signal, _) = watch::channel(false);
160        let (progress, _) = watch::channel(Ok(origin));
161        let shared = Arc::new(Shared { failed, signal, progress, capacity: Arc::new(Semaphore::new(self.limits.pending_bytes as usize)) });
162        let mut guard = Failure { shared: shared.clone(), clean: false };
163        let mut consumer = self.consumer;
164        let worker = tokio::spawn(async move {
165            while let Some(mut group) = receiver.recv().await {
166                group.trace.span.follows_from(group.trace.predecessor());
167                group.trace.dequeued();
168                if consumer
169                    .capture(group.entries, group.frontier)
170                    .instrument(group.trace.span.clone())
171                    .in_context(group.trace.dispatcher.clone())
172                    .await
173                    .is_err()
174                {
175                    return;
176                }
177                // Capacity includes the callback until successful acknowledgement.
178                let _ = guard.shared.progress.send_replace(Ok(group.frontier));
179                drop(group._permit);
180            }
181            guard.clean = true;
182            drop(guard);
183        });
184        Callback { sender, shared, limits: self.limits, worker }
185    }
186}
187
188impl<A: Application> Callback<A> {
189    #[tracing::instrument(target = "harmos::state", name = "state.reserve", level = "debug", skip_all)]
190    pub async fn reserve(&self, bytes: usize) -> Result<OwnedSemaphorePermit, PreparationError> {
191        self.check()?;
192        if bytes > self.limits.max_group_bytes as usize {
193            return Err(PreparationError::Limit);
194        }
195        self.shared.capacity.clone().acquire_many_owned(bytes as u32).await.map_err(|_| PreparationError::CaptureFailed)
196    }
197
198    pub fn check(&self) -> Result<(), PreparationError> {
199        if self.shared.failed.load(Ordering::Acquire) { Err(PreparationError::CaptureFailed) } else { Ok(()) }
200    }
201
202    pub fn submit(&self, entries: Vec<Arc<Entry<Recorded<A>>>>, frontier: Position, permit: OwnedSemaphorePermit) {
203        if self
204            .sender
205            .send(Group {
206                entries,
207                frontier,
208                _permit: permit,
209                trace: crate::supervision::executor::SpanContext::new(
210                    tracing::debug_span!(target: "harmos::state", "state.capture_group", frontier = frontier.0),
211                    tracing::debug_span!(target: "harmos::state", "state.capture_queue"),
212                ),
213            })
214            .is_err()
215        {
216            self.shared.fail();
217        }
218    }
219
220    pub fn progress(&self) -> watch::Receiver<Result<Position, CaptureError>> {
221        self.shared.progress.subscribe()
222    }
223
224    pub fn failure(&self) -> watch::Receiver<bool> {
225        self.shared.signal.subscribe()
226    }
227
228    #[tracing::instrument(target = "harmos::state", name = "state.finish", level = "debug", skip_all)]
229    pub async fn finish(self) {
230        drop(self.sender);
231        let _ = self.worker.await;
232    }
233}
234
235impl<A: Application> Declaration<A> {
236    pub(crate) fn start(self, failed: Arc<AtomicBool>, origin: Position) -> Capture<A> {
237        match self {
238            Self::Callback(callback) => Capture::Callback(callback.start(failed, origin)),
239            Self::Managed { storage, receipts, projection, snapshots } => Capture::Managed {
240                storage,
241                receipts,
242                projection,
243                snapshots,
244                failed,
245                progress: watch::channel(Ok(origin)).0,
246                signal: watch::channel(false).0,
247            },
248        }
249    }
250}
251
252impl<A: Application> Capture<A> {
253    #[tracing::instrument(target = "harmos::state", name = "state.snapshot", level = "debug", skip_all)]
254    pub(crate) async fn snapshot(
255        &self, tracked: &super::TrackedState<A>, undos: &super::Undos<A>, receipts: &HashMap<RequestKey, Receipt>, final_save: bool,
256    ) {
257        if let Self::Managed { storage, snapshots: Some(snapshots), .. } = self {
258            snapshots.capture(storage, tracked, undos, receipts, final_save).await;
259        }
260    }
261
262    pub(crate) fn take_receipts(&mut self) -> HashMap<RequestKey, Receipt> {
263        match self {
264            Self::Callback(_) => HashMap::new(),
265            Self::Managed { receipts, .. } => std::mem::take(receipts),
266        }
267    }
268
269    #[tracing::instrument(target = "harmos::state", name = "state.reserve", level = "debug", skip_all)]
270    pub(crate) async fn reserve(
271        &self, bytes: usize, group: &crate::state::recovery::Group<&serde_json::Value>,
272    ) -> Result<Permit, PreparationError> {
273        match self {
274            Self::Callback(callback) => callback.reserve(bytes).await.map(Permit::Callback),
275            Self::Managed { storage, projection, .. } => {
276                storage.reserve_group(group, projection.as_ref()).await.map(|permit| Permit::Managed(Box::new(permit))).map_err(|error| {
277                    match error {
278                        crate::Error::InvalidInput { .. } => PreparationError::Limit,
279                        _ => {
280                            self.fail();
281                            PreparationError::CaptureFailed
282                        }
283                    }
284                })
285            }
286        }
287    }
288
289    pub(crate) fn check(&self) -> Result<(), PreparationError> {
290        match self {
291            Self::Callback(callback) => callback.check(),
292            Self::Managed { storage, .. } => {
293                let status = storage.status();
294                if status.failure.is_some() || status.closing || !status.running {
295                    self.fail();
296                    Err(PreparationError::CaptureFailed)
297                } else {
298                    Ok(())
299                }
300            }
301        }
302    }
303
304    pub(crate) fn submit(&self, entries: Vec<Arc<Entry<Recorded<A>>>>, frontier: Position, permit: Permit) {
305        match (self, permit) {
306            (Self::Callback(callback), Permit::Callback(permit)) => callback.submit(entries, frontier, permit),
307            (Self::Managed { progress, .. }, Permit::Managed(permit)) => {
308                (*permit).commit(frontier);
309                let _ = progress.send_replace(Ok(frontier));
310            }
311            _ => unreachable!("capture owns its reservation"),
312        }
313    }
314
315    pub(crate) fn progress(&self) -> watch::Receiver<Result<Position, CaptureError>> {
316        match self {
317            Self::Callback(callback) => callback.progress(),
318            Self::Managed { progress, .. } => progress.subscribe(),
319        }
320    }
321
322    pub(crate) fn failure(&self) -> watch::Receiver<bool> {
323        match self {
324            Self::Callback(callback) => callback.failure(),
325            Self::Managed { signal, .. } => signal.subscribe(),
326        }
327    }
328
329    fn fail(&self) {
330        if let Self::Managed { failed, progress, signal, .. } = self {
331            failed.store(true, Ordering::Release);
332            let _ = progress.send_replace(Err(CaptureError::Failed));
333            signal.send_replace(true);
334        }
335    }
336
337    #[tracing::instrument(target = "harmos::state", name = "state.finish", level = "debug", skip_all)]
338    pub(crate) async fn finish(self) {
339        if let Self::Callback(callback) = self {
340            callback.finish().await;
341        }
342    }
343}
344
345/* === Public Functions === */
346
347/* === Private Functions === */
348
349/* === Tests === */
harmos_runtime/state/capture.rs

harmos_runtime/state/
capture.rs
