harmos_storage/
execution.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::future::Future;
18use std::pin::Pin;
19use std::task::{Context as TaskContext, Poll};
20use std::time::Duration;
21
22// Third Party Imports
23use tokio::time::Instant;
24
25// Local Imports
26
27/* === Constants === */
28
29/* === Macros === */
30
31/* === Type Aliases === */
32
33pub type Sleep = Pin<Box<dyn Future<Output = ()> + Send>>;
34
35/* === Structs === */
36
37#[derive(Clone, Copy, Debug, PartialEq, Eq)]
38pub struct Elapsed;
39
40#[derive(Clone, Copy)]
41pub struct Execution {
42    now: fn() -> Instant,
43    sleep: fn(Duration) -> Sleep,
44}
45
46#[derive(Clone)]
47pub(crate) struct Context {
48    pub(crate) span: tracing::Span,
49    dispatch: tracing::Dispatch,
50}
51
52struct Scoped<F> {
53    context: Context,
54    future: Option<Pin<Box<F>>>,
55}
56
57/* === Enums === */
58
59/* === Traits === */
60
61/* === Implementations === */
62
63impl Execution {
64    pub const fn new(now: fn() -> Instant, sleep: fn(Duration) -> Sleep) -> Self {
65        Self { now, sleep }
66    }
67    pub fn now(self) -> Instant {
68        (self.now)()
69    }
70    pub fn sleep(self, duration: Duration) -> Sleep {
71        (self.sleep)(duration)
72    }
73    #[tracing::instrument(name = "storage.execution.timeout", level = "debug", skip_all)]
74    pub async fn timeout<F: Future>(self, duration: Duration, future: F) -> Result<F::Output, Elapsed> {
75        tokio::select! { biased; result = future => Ok(result), () = self.sleep(duration) => Err(Elapsed) }
76    }
77}
78
79impl Context {
80    pub(crate) fn current() -> Self {
81        Self { span: tracing::Span::current(), dispatch: tracing::dispatcher::get_default(Clone::clone) }
82    }
83
84    pub(crate) fn with_span(span: tracing::Span) -> Self {
85        Self { span, dispatch: tracing::dispatcher::get_default(Clone::clone) }
86    }
87
88    pub(crate) async fn scope<F: Future>(&self, future: F) -> F::Output {
89        Scoped { context: self.clone(), future: Some(Box::pin(future)) }.await
90    }
91
92    pub(crate) async fn result<T>(&self, span: tracing::Span, future: impl Future<Output = crate::Result<T>>) -> crate::Result<T> {
93        let future = async {
94            let result = future.await;
95            tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
96            result
97        };
98        self.scope(tracing::Instrument::instrument(future, span)).await
99    }
100
101    pub(crate) fn sync<T>(&self, action: impl FnOnce() -> T) -> T {
102        tracing::dispatcher::with_default(&self.dispatch, || self.span.in_scope(action))
103    }
104
105    pub(crate) fn blocking<T>(&self, action: impl FnOnce() -> crate::Result<T>) -> crate::Result<T> {
106        self.sync(|| {
107            let span = tracing::debug_span!("storage.backend.blocking", outcome = "cancelled");
108            span.in_scope(|| {
109                let result = action();
110                span.record("outcome", if result.is_ok() { "completed" } else { "failed" });
111                result
112            })
113        })
114    }
115}
116
117impl Drop for Context {
118    fn drop(&mut self) {
119        let span = std::mem::replace(&mut self.span, tracing::Span::none());
120        tracing::dispatcher::with_default(&self.dispatch, || drop(span));
121    }
122}
123
124impl<F: Future> Future for Scoped<F> {
125    type Output = F::Output;
126
127    fn poll(self: Pin<&mut Self>, task: &mut TaskContext<'_>) -> Poll<Self::Output> {
128        let this = self.get_mut();
129        let context = &this.context;
130        let future = this.future.as_mut().expect("scoped future remains owned while polling");
131        context.sync(|| future.as_mut().poll(task))
132    }
133}
134
135impl<F> Drop for Scoped<F> {
136    fn drop(&mut self) {
137        let future = self.future.take();
138        self.context.sync(|| drop(future));
139    }
140}
141
142impl Default for Execution {
143    fn default() -> Self {
144        Self::new(Instant::now, |duration| Box::pin(tokio::time::sleep(duration)))
145    }
146}
147
148/* === Public Functions === */
149
150/* === Private Functions === */
151
152/* === Tests === */
harmos_storage/execution.rs

harmos_storage/
execution.rs
