coven_storage/cloud/
runtime.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::{Arc, Mutex};
4
5use super::{CloudFileReadError, CloudHomeError};
6
7tokio::task_local! {
8    static CLOUD_RUNTIME_TASK: ();
9}
10
11type CloudFuture<T, E> = Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'static>>;
12type CloudOperation<T, E> = Box<dyn FnOnce() -> CloudFuture<T, E> + Send + 'static>;
13type CloudRun<T, E> =
14    Pin<Box<dyn Future<Output = Result<Result<T, E>, CloudRuntimeError>> + Send + 'static>>;
15type CloudTask<T> = Pin<Box<dyn Future<Output = T> + Send + 'static>>;
16type CloudTaskOperation<T> = Box<dyn FnOnce() -> CloudTask<T> + Send + 'static>;
17
18/// A failure to start or execute work on Coven's cloud runtime.
19#[derive(Debug, thiserror::Error)]
20pub enum CloudRuntimeError {
21    #[error("start Coven cloud runtime: {0}")]
22    Start(#[source] std::io::Error),
23    #[error("cloud operation task failed: {0}")]
24    Task(#[source] tokio::task::JoinError),
25}
26
27/// The executor retained by Coven's cloud owner and the provider homes it
28/// creates. The Tokio runtime is built on first use, and callers provide an
29/// operation builder so its future is also constructed on the cloud worker.
30#[derive(Clone)]
31pub(crate) struct CloudRuntime {
32    inner: Arc<CloudRuntimeInner>,
33}
34
35struct CloudRuntimeInner {
36    runtime: Mutex<Option<tokio::runtime::Runtime>>,
37}
38
39impl CloudRuntime {
40    pub(crate) fn new() -> Self {
41        Self {
42            inner: Arc::new(CloudRuntimeInner {
43                runtime: Mutex::new(None),
44            }),
45        }
46    }
47
48    pub(crate) fn run<T, E, F>(
49        &self,
50        operation: impl FnOnce() -> F + Send + 'static,
51    ) -> CloudRun<T, E>
52    where
53        T: Send + 'static,
54        E: Send + 'static,
55        F: Future<Output = Result<T, E>> + Send + 'static,
56    {
57        self.run_erased(Box::new(move || Box::pin(operation())))
58    }
59
60    fn run_erased<T, E>(&self, operation: CloudOperation<T, E>) -> CloudRun<T, E>
61    where
62        T: Send + 'static,
63        E: Send + 'static,
64    {
65        let runtime = self.clone();
66        Box::pin(async move {
67            if CLOUD_RUNTIME_TASK.try_with(|_| ()).is_ok() {
68                return Ok(operation().await);
69            }
70
71            AbortOnDropTask::new(runtime.spawn_erased(operation)?)
72                .wait()
73                .await
74                .map_err(CloudRuntimeError::Task)
75        })
76    }
77
78    pub(crate) async fn run_cloud<T, F>(
79        &self,
80        operation: impl FnOnce() -> F + Send + 'static,
81    ) -> Result<T, CloudHomeError>
82    where
83        T: Send + 'static,
84        F: Future<Output = Result<T, CloudHomeError>> + Send + 'static,
85    {
86        self.run(operation)
87            .await
88            .map_err(|error| CloudHomeError::transport("run cloud operation", error))?
89    }
90
91    pub(crate) async fn run_file_read<T, F>(
92        &self,
93        operation: impl FnOnce() -> F + Send + 'static,
94    ) -> Result<T, CloudFileReadError>
95    where
96        T: Send + 'static,
97        F: Future<Output = Result<T, CloudFileReadError>> + Send + 'static,
98    {
99        self.run(operation).await.map_err(|error| {
100            CloudFileReadError::Source(CloudHomeError::transport(
101                "run cloud file-read operation",
102                error,
103            ))
104        })?
105    }
106
107    pub(crate) fn spawn<T, F>(
108        &self,
109        operation: impl FnOnce() -> F + Send + 'static,
110    ) -> Result<tokio::task::JoinHandle<T>, CloudRuntimeError>
111    where
112        T: Send + 'static,
113        F: Future<Output = T> + Send + 'static,
114    {
115        self.spawn_erased(Box::new(move || Box::pin(operation())))
116    }
117
118    fn spawn_erased<T>(
119        &self,
120        operation: CloudTaskOperation<T>,
121    ) -> Result<tokio::task::JoinHandle<T>, CloudRuntimeError>
122    where
123        T: Send + 'static,
124    {
125        let handle = self.handle()?;
126        let lifetime = self.clone();
127        Ok(handle.spawn(CLOUD_RUNTIME_TASK.scope((), async move {
128            let result = operation().await;
129            drop(lifetime);
130            result
131        })))
132    }
133
134    fn handle(&self) -> Result<tokio::runtime::Handle, CloudRuntimeError> {
135        let mut runtime = self.inner.runtime.lock().expect("lock cloud runtime");
136        if runtime.is_none() {
137            *runtime = Some(
138                tokio::runtime::Builder::new_multi_thread()
139                    .worker_threads(1)
140                    .thread_stack_size(16 * 1024 * 1024)
141                    .thread_name("coven-cloud")
142                    .enable_all()
143                    .build()
144                    .map_err(CloudRuntimeError::Start)?,
145            );
146        }
147        Ok(runtime
148            .as_ref()
149            .expect("cloud runtime initialized above")
150            .handle()
151            .clone())
152    }
153}
154
155impl Drop for CloudRuntimeInner {
156    fn drop(&mut self) {
157        if let Some(runtime) = self
158            .runtime
159            .get_mut()
160            .expect("lock cloud runtime during final drop")
161            .take()
162        {
163            runtime.shutdown_background();
164        }
165    }
166}
167
168struct AbortOnDropTask<T> {
169    handle: Option<tokio::task::JoinHandle<T>>,
170}
171
172impl<T> AbortOnDropTask<T> {
173    fn new(handle: tokio::task::JoinHandle<T>) -> Self {
174        Self {
175            handle: Some(handle),
176        }
177    }
178
179    async fn wait(mut self) -> Result<T, tokio::task::JoinError> {
180        let result = self
181            .handle
182            .as_mut()
183            .expect("cloud task handle is present")
184            .await;
185        self.handle.take();
186        result
187    }
188}
189
190impl<T> Drop for AbortOnDropTask<T> {
191    fn drop(&mut self) {
192        if let Some(handle) = self.handle.take() {
193            handle.abort();
194        }
195    }
196}
197
198#[cfg(test)]
199mod tests;
coven_storage/cloud/runtime.rs

coven_storage/cloud/
runtime.rs
