fedimint_core/task/
jit.rs

1use std::convert::Infallible;
2use std::fmt;
3#[cfg(not(target_family = "wasm"))]
4use std::panic;
5use std::sync::Arc;
6
7use fedimint_core::runtime::JoinHandle;
8use fedimint_logging::LOG_TASK;
9use futures::Future;
10use tokio::sync;
11use tracing::warn;
12
13use super::MaybeSend;
14use crate::util::FmtCompact;
15
16pub type Jit<T> = JitCore<T, Infallible>;
17pub type JitTry<T, E> = JitCore<T, E>;
18
19/// Error that could have been returned before
20///
21/// Newtype over `Option<E>` that allows better user (error conversion mostly)
22/// experience
23#[derive(Debug)]
24pub enum OneTimeError<E> {
25    Original(E),
26    Copy(String),
27}
28
29impl<E> std::error::Error for OneTimeError<E>
30where
31    E: fmt::Debug + fmt::Display,
32{
33    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
34        // In neither case we can preserve the source information
35        None
36    }
37
38    fn cause(&self) -> Option<&dyn std::error::Error> {
39        self.source()
40    }
41}
42
43impl<E> fmt::Display for OneTimeError<E>
44where
45    E: fmt::Display,
46{
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        match self {
49            Self::Original(o) => o.fmt(f),
50            Self::Copy(c) => f.write_str(c),
51        }
52    }
53}
54
55/// A value that initializes eagerly in parallel in a falliable way
56#[derive(Debug)]
57pub struct JitCore<T, E> {
58    inner: Arc<JitInner<T, E>>,
59}
60
61#[derive(Debug)]
62struct JitInner<T, E> {
63    handle: sync::Mutex<JoinHandle<Result<T, E>>>,
64    val: sync::OnceCell<Result<T, String>>,
65}
66
67impl<T, E> Clone for JitCore<T, E>
68where
69    T: Clone,
70{
71    fn clone(&self) -> Self {
72        Self {
73            inner: self.inner.clone(),
74        }
75    }
76}
77impl<T, E> Drop for JitInner<T, E> {
78    fn drop(&mut self) {
79        self.handle.get_mut().abort();
80    }
81}
82impl<T, E> JitCore<T, E>
83where
84    T: MaybeSend + 'static,
85    E: MaybeSend + 'static + fmt::Display,
86{
87    /// Create `JitTry` value, and spawn a future `f` that computes its value
88    ///
89    /// Unlike normal Rust futures, the `f` executes eagerly (is spawned as a
90    /// tokio task).
91    pub fn new_try<Fut>(f: impl FnOnce() -> Fut + 'static + MaybeSend) -> Self
92    where
93        Fut: Future<Output = std::result::Result<T, E>> + 'static + MaybeSend,
94    {
95        let handle = crate::runtime::spawn("jit-value", async { f().await });
96
97        Self {
98            inner: JitInner {
99                handle: handle.into(),
100                val: sync::OnceCell::new(),
101            }
102            .into(),
103        }
104    }
105
106    /// Get the reference to the value, potentially blocking for the
107    /// initialization future to complete
108    pub async fn get_try(&self) -> Result<&T, OneTimeError<E>> {
109        let mut init_error = None;
110        let value = self
111            .inner
112            .val
113            .get_or_init(|| async {
114                let handle: &mut _ = &mut *self.inner.handle.lock().await;
115                match handle.await {
116                        Ok(Ok(o)) => Ok(o),
117                        Ok(Err(err)) => {
118                            let err_str = err.to_string();
119                            init_error = Some(err);
120                            Err(err_str)
121                        },
122                        Err(err) => {
123
124                            #[cfg(not(target_family = "wasm"))]
125                            if err.is_panic() {
126                                warn!(target: LOG_TASK, err = %err.fmt_compact(), type_name = %std::any::type_name::<T>(), "Jit value panicked");
127                                // Resume the panic on the main task
128                                panic::resume_unwind(err.into_panic());
129                            }
130                            #[cfg(not(target_family = "wasm"))]
131                            if err.is_cancelled() {
132                                warn!(target: LOG_TASK, err = %err.fmt_compact(), type_name = %std::any::type_name::<T>(), "Jit value task canceled:");
133                            }
134                            Err(format!("Jit value {} failed unexpectedly with: {}", std::any::type_name::<T>(), err))
135                        },
136                    }
137            })
138            .await;
139        if let Some(err) = init_error {
140            return Err(OneTimeError::Original(err));
141        }
142        value
143            .as_ref()
144            .map_err(|err_str| OneTimeError::Copy(err_str.to_owned()))
145    }
146}
147impl<T> JitCore<T, Infallible>
148where
149    T: MaybeSend + 'static,
150{
151    pub fn new<Fut>(f: impl FnOnce() -> Fut + 'static + MaybeSend) -> Self
152    where
153        Fut: Future<Output = T> + 'static + MaybeSend,
154        T: 'static,
155    {
156        Self::new_try(|| async { Ok(f().await) })
157    }
158
159    pub async fn get(&self) -> &T {
160        self.get_try().await.expect("can't fail")
161    }
162}
163#[cfg(test)]
164mod tests;
fedimint_core/task/jit.rs

fedimint_core/task/
jit.rs
