saluki_common/task/
mod.rs

1//! Helpers for working with asynchronous tasks.
2
3use std::future::Future;
4
5use tokio::{
6    runtime::Handle,
7    task::{AbortHandle, JoinHandle, JoinSet},
8};
9use tracing::Instrument as _;
10
11use crate::resource_tracking::Track as _;
12
13mod instrument;
14use self::instrument::TaskInstrument as _;
15
16/// Spawns a new asynchronous task, returning a [`JoinHandle`] for it.
17///
18/// This function is a thin wrapper over [`tokio::spawn`], and provides implicit "tracing" for spawned futures by
19/// ensuring that the task is attached to the current `tracing` span and the current allocation component.
20#[track_caller]
21pub fn spawn_traced<F, T>(f: F) -> JoinHandle<T>
22where
23    F: Future<Output = T> + Send + 'static,
24    T: Send + 'static,
25{
26    tokio::task::spawn(
27        f.in_current_span()
28            .in_current_resource_group()
29            .with_task_instrumentation(get_caller_location_as_string()),
30    )
31}
32
33/// Spawns a new named asynchronous task, returning a [`JoinHandle`] for it.
34///
35/// This function is a thin wrapper over [`tokio::spawn`], and provides implicit "tracing" for spawned futures by
36/// ensuring that the task is attached to the current `tracing` span and the current allocation component.
37pub fn spawn_traced_named<S, F, T>(name: S, f: F) -> JoinHandle<T>
38where
39    S: Into<String>,
40    F: Future<Output = T> + Send + 'static,
41    T: Send + 'static,
42{
43    tokio::task::spawn(
44        f.in_current_span()
45            .in_current_resource_group()
46            .with_task_instrumentation(name.into()),
47    )
48}
49
50/// Helper trait for providing traced spawning when using `JoinSet<T>`.
51pub trait JoinSetExt<T> {
52    /// Spawns a new asynchronous task, returning an [`AbortHandle`] for it.
53    ///
54    /// This is meant to be a thin wrapper over task management types like [`JoinSet`], and provides implicit "tracing"
55    /// for spawned futures by ensuring that the task is attached to the current `tracing` span and the current
56    /// allocation component.
57    fn spawn_traced<F>(&mut self, f: F) -> AbortHandle
58    where
59        F: Future<Output = T> + Send + 'static,
60        T: Send + 'static;
61
62    /// Spawns a new named asynchronous task, returning an [`AbortHandle`] for it.
63    ///
64    /// This is meant to be a thin wrapper over task management types like [`JoinSet`], and provides implicit "tracing"
65    /// for spawned futures by ensuring that the task is attached to the current `tracing` span and the current
66    /// allocation component.
67    fn spawn_traced_named<S, F>(&mut self, name: S, f: F) -> AbortHandle
68    where
69        S: Into<String>,
70        F: Future<Output = T> + Send + 'static,
71        T: Send + 'static;
72}
73
74impl<T> JoinSetExt<T> for JoinSet<T> {
75    fn spawn_traced<F>(&mut self, f: F) -> AbortHandle
76    where
77        F: Future<Output = T> + Send + 'static,
78        T: Send + 'static,
79    {
80        self.spawn(
81            f.in_current_span()
82                .in_current_resource_group()
83                .with_task_instrumentation(get_caller_location_as_string()),
84        )
85    }
86
87    fn spawn_traced_named<S, F>(&mut self, name: S, f: F) -> AbortHandle
88    where
89        S: Into<String>,
90        F: Future<Output = T> + Send + 'static,
91        T: Send + 'static,
92    {
93        self.spawn(
94            f.in_current_span()
95                .in_current_resource_group()
96                .with_task_instrumentation(name.into()),
97        )
98    }
99}
100
101/// Helper trait for providing traced spawning when using `Handle`.
102pub trait HandleExt<T> {
103    /// Spawns a new asynchronous task, returning a [`JoinHandle`] for it.
104    ///
105    /// This is meant to be a thin wrapper over task management types like [`Handle`], and provides implicit "tracing"
106    /// for spawned futures by ensuring that the task is attached to the current `tracing` span and the current
107    /// allocation component.
108    fn spawn_traced<F>(&self, f: F) -> JoinHandle<T>
109    where
110        F: Future<Output = T> + Send + 'static,
111        T: Send + 'static;
112
113    /// Spawns a new named asynchronous task, returning a [`JoinHandle`] for it.
114    ///
115    /// This is meant to be a thin wrapper over task management types like [`Handle`], and provides implicit "tracing"
116    /// for spawned futures by ensuring that the task is attached to the current `tracing` span and the current
117    /// allocation component.
118    fn spawn_traced_named<S, F>(&self, name: S, f: F) -> JoinHandle<T>
119    where
120        S: Into<String>,
121        F: Future<Output = T> + Send + 'static,
122        T: Send + 'static;
123}
124
125impl<T> HandleExt<T> for Handle {
126    fn spawn_traced<F>(&self, f: F) -> JoinHandle<T>
127    where
128        F: Future<Output = T> + Send + 'static,
129        T: Send + 'static,
130    {
131        self.spawn(
132            f.in_current_span()
133                .in_current_resource_group()
134                .with_task_instrumentation(get_caller_location_as_string()),
135        )
136    }
137
138    fn spawn_traced_named<S, F>(&self, name: S, f: F) -> JoinHandle<T>
139    where
140        S: Into<String>,
141        F: Future<Output = T> + Send + 'static,
142        T: Send + 'static,
143    {
144        self.spawn(
145            f.in_current_span()
146                .in_current_resource_group()
147                .with_task_instrumentation(name.into()),
148        )
149    }
150}
151
152/// Gets the caller location as a string, in the form of `file:line:column`.
153#[track_caller]
154pub fn get_caller_location_as_string() -> String {
155    let caller = std::panic::Location::caller();
156    format!("file-{}@{}-{}", caller.file(), caller.line(), caller.column())
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn caller_location_uses_documented_format() {
165        // Because `get_caller_location_as_string` is `#[track_caller]`, this records _this_ call
166        // site, so the file segment matches this source file.
167        let location = get_caller_location_as_string();
168
169        // NOTE: the doc comment describes the format as `file:line:column`, but the implementation
170        // actually produces `file-<file>@<line>-<column>`. We assert the real, current format.
171        let prefix = format!("file-{}@", file!());
172        assert!(
173            location.starts_with(&prefix),
174            "expected {location:?} to start with {prefix:?}"
175        );
176
177        // Exactly one `@` separates the file from the position, and the trailing portion is two
178        // decimal numbers (line and column) joined by `-`.
179        assert_eq!(location.matches('@').count(), 1);
180        let (_, line_col) = location.split_once('@').unwrap();
181        let (line, column) = line_col.split_once('-').expect("line and column separated by '-'");
182        assert!(line.parse::<u32>().is_ok(), "line segment {line:?} should be numeric");
183        assert!(
184            column.parse::<u32>().is_ok(),
185            "column segment {column:?} should be numeric"
186        );
187    }
188
189    #[tokio::test]
190    async fn spawn_traced_runs_future_to_completion() {
191        assert_eq!(spawn_traced(async { 7u32 }).await.unwrap(), 7);
192        assert_eq!(spawn_traced_named("named-task", async { 9u32 }).await.unwrap(), 9);
193    }
194
195    #[tokio::test]
196    async fn join_set_ext_runs_traced_tasks_to_completion() {
197        let mut set = JoinSet::new();
198        set.spawn_traced(async { 1u32 });
199        set.spawn_traced_named("set-task", async { 2u32 });
200
201        let mut results = Vec::new();
202        while let Some(result) = set.join_next().await {
203            results.push(result.unwrap());
204        }
205        results.sort_unstable();
206        assert_eq!(results, vec![1, 2]);
207    }
208
209    #[tokio::test]
210    async fn handle_ext_runs_traced_tasks_to_completion() {
211        let handle = Handle::current();
212        assert_eq!(handle.spawn_traced(async { 4u32 }).await.unwrap(), 4);
213        assert_eq!(
214            handle.spawn_traced_named("handle-task", async { 5u32 }).await.unwrap(),
215            5
216        );
217    }
218}
saluki_common/task/mod.rs

saluki_common/task/
mod.rs
