harmos_sidecar/kit/
task.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, io, sync::mpsc, thread};
18// Third Party Imports
19use tokio::{runtime::Builder, sync::oneshot, task::JoinHandle};
20use tracing::Instrument;
21// Local Imports
22
23/* === Constants === */
24
25/* === Macros === */
26
27/* === Type Aliases === */
28
29/* === Structs === */
30
31struct Contextual<F: std::future::Future> {
32    future: Option<std::pin::Pin<Box<F>>>,
33    dispatcher: tracing::Dispatch,
34}
35
36/* === Enums === */
37
38/* === Traits === */
39
40trait InContext: std::future::Future + Sized {
41    fn in_context(self, dispatcher: tracing::Dispatch) -> Contextual<Self> {
42        Contextual { future: Some(Box::pin(self)), dispatcher }
43    }
44
45    fn in_current_context(self) -> Contextual<Self> {
46        self.in_context(tracing::dispatcher::get_default(Clone::clone))
47    }
48}
49
50/* === Implementations === */
51
52impl<F: std::future::Future> InContext for F {}
53
54impl<F: std::future::Future> std::future::Future for Contextual<F> {
55    type Output = F::Output;
56
57    fn poll(self: std::pin::Pin<&mut Self>, context: &mut std::task::Context<'_>) -> std::task::Poll<Self::Output> {
58        let owner = self.get_mut();
59        let _dispatcher = tracing::dispatcher::set_default(&owner.dispatcher);
60        owner.future.as_mut().expect("context future remains until drop").as_mut().poll(context)
61    }
62}
63
64impl<F: std::future::Future> Drop for Contextual<F> {
65    fn drop(&mut self) {
66        tracing::dispatcher::with_default(&self.dispatcher, || drop(self.future.take()));
67    }
68}
69
70/* === Public Functions === */
71
72pub fn spawn(name: String, work: impl Future<Output = ()> + Send + 'static) -> io::Result<JoinHandle<()>> {
73    let work = work.instrument(tracing::Span::current()).in_current_context();
74    let (started, ready) = mpsc::sync_channel(1);
75    thread::Builder::new().name(name).spawn(move || {
76        let runtime = match Builder::new_current_thread().enable_all().build() {
77            Ok(runtime) => runtime,
78            Err(error) => {
79                let _ = started.send(Err(error));
80                return;
81            }
82        };
83        let (finished, completion) = oneshot::channel::<()>();
84        let task = runtime.spawn(async move {
85            // Dropping this guard also wakes the runtime owner after an abort
86            // or panic. It must live until the author's future has dropped.
87            let _finished = finished;
88            work.await;
89        });
90        if started.send(Ok(task)).is_err() {
91            return;
92        }
93
94        let _ = runtime.block_on(completion);
95    })?;
96
97    ready.recv().map_err(io::Error::other)?
98}
99
100/* === Private Functions === */
101
102/* === Tests === */
harmos_sidecar/kit/task.rs

harmos_sidecar/kit/
task.rs
