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 {
82            span: tracing::Span::current(),
83            dispatch: tracing::dispatcher::get_default(Clone::clone),
84        }
85    }
86
87    pub(crate) fn with_span(span: tracing::Span) -> Self {
88        Self {
89            span,
90            dispatch: tracing::dispatcher::get_default(Clone::clone),
91        }
92    }
93
94    pub(crate) async fn scope<F: Future>(&self, future: F) -> F::Output {
95        Scoped {
96            context: self.clone(),
97            future: Some(Box::pin(future)),
98        }
99        .await
100    }
101
102    pub(crate) async fn result<T>(&self, span: tracing::Span, future: impl Future<Output = crate::Result<T>>) -> crate::Result<T> {
103        let future = async {
104            let result = future.await;
105            tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
106            result
107        };
108        self.scope(tracing::Instrument::instrument(future, span)).await
109    }
110
111    pub(crate) fn sync<T>(&self, action: impl FnOnce() -> T) -> T {
112        tracing::dispatcher::with_default(&self.dispatch, || self.span.in_scope(action))
113    }
114
115    pub(crate) fn blocking<T>(&self, action: impl FnOnce() -> crate::Result<T>) -> crate::Result<T> {
116        self.sync(|| {
117            let span = tracing::debug_span!("storage.backend.blocking", outcome = "cancelled");
118            span.in_scope(|| {
119                let result = action();
120                span.record("outcome", if result.is_ok() { "completed" } else { "failed" });
121                result
122            })
123        })
124    }
125}
126
127impl Drop for Context {
128    fn drop(&mut self) {
129        let span = std::mem::replace(&mut self.span, tracing::Span::none());
130        tracing::dispatcher::with_default(&self.dispatch, || drop(span));
131    }
132}
133
134impl<F: Future> Future for Scoped<F> {
135    type Output = F::Output;
136
137    fn poll(self: Pin<&mut Self>, task: &mut TaskContext<'_>) -> Poll<Self::Output> {
138        let this = self.get_mut();
139        let context = &this.context;
140        let future = this.future.as_mut().expect("scoped future remains owned while polling");
141        context.sync(|| future.as_mut().poll(task))
142    }
143}
144
145impl<F> Drop for Scoped<F> {
146    fn drop(&mut self) {
147        let future = self.future.take();
148        self.context.sync(|| drop(future));
149    }
150}
151
152impl Default for Execution {
153    fn default() -> Self {
154        Self::new(Instant::now, |duration| Box::pin(tokio::time::sleep(duration)))
155    }
156}
157
158/* === Public Functions === */
159
160/* === Private Functions === */
161
162/* === Tests === */
harmos_storage/execution.rs

harmos_storage/
execution.rs
