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}