tracing_tools/
lib.rs

1use std::{pin::Pin, time::Instant};
2use std::future::Future;
3
4use tracing::{info, error, span::{Span}};
5use tracing_futures::Instrument;
6
7type Result<T> = anyhow::Result<T>;
8pub type TaskFut<'a, T=()> = Pin<Box<dyn Future<Output=Result<T>> + 'a>>;
9
10pub struct TracingTask<'a, R=()> {
11    span: Span,
12    future: TaskFut<'a, R>,
13    is_long_lived: bool
14}
15
16impl<'a, R> TracingTask<'a, R> {
17    pub fn new<T: Future<Output=Result<R>> + 'a>(span: Span, fut: T) -> TracingTask<'a, R> {
18        TracingTask {
19            span,
20            future: Box::pin(fut),
21            is_long_lived: true
22        }
23    }
24
25    pub fn new_short_lived<T: Future<Output=Result<R>> + 'a>(span: Span, fut: T) -> TracingTask<'a, R> {
26        TracingTask {
27            span,
28            future: Box::pin(fut),
29            is_long_lived: false
30        }
31    }
32}
33
34impl<'a, R: 'a> TracingTask<'a, R> {
35    pub fn instrument(self) -> TaskFut<'a, R> {
36        let span = self.span;
37        let future = self.future;
38        let is_long_lived = self.is_long_lived;
39
40        let fut_wrap = async move {
41            if is_long_lived {
42                info!("Starting...");
43            }
44            let t = Instant::now();
45
46            let r = future.await;
47            if r.is_err() {
48                let err = r.err().unwrap();
49                error!(error = ?err, elapsed = ?t.elapsed(), "Finished with");
50                return Err(err);
51            }
52            info!(elapsed = ?t.elapsed(), "Finished [OK]...");
53            Ok(r.unwrap())
54        };
55
56        Box::pin(fut_wrap.instrument(span))
57    }
58}
59
60pub fn clean_fn(s: &str) -> String {
61    let s = String::from(s);
62    let name = s.split("::")
63        .collect::<Vec<&str>>()
64        .into_iter().rev()
65        .take(2).rev()
66        .collect::<Vec<&str>>()
67        .join("::");
68
69    let mut final_name = String::from("");
70    let mut skip = 0;
71    for c in name.chars() {
72        if c == '<' {
73            skip += 1;
74        } else if c == '>' {
75            skip -= 1;
76        } else if skip < 1 {
77            final_name.push(c);
78        }
79    }
80    final_name
81}
82
83#[macro_export]
84macro_rules! function {
85    () => {{
86        fn f() {}
87        fn type_name_of<T>(_: T) -> &'static str {
88            std::any::type_name::<T>()
89        }
90        let name = type_name_of(f);
91        &name[..name.len() - 3]
92    }}
93}
94
95#[macro_export]
96macro_rules! span {
97    ($($tts:tt)*) => {
98        tracing::span!(tracing::Level::ERROR, "task", name = $crate::clean_fn($crate::function!()).as_str(), $($tts)*);
99    };
100    ($name:expr) => {
101        tracing::span!(tracing::Level::ERROR, "task", name = $name);
102    };
103    () => {
104        tracing::span!(tracing::Level::ERROR, "task", name = $crate::clean_fn($crate::function!()).as_str());
105    };
106}