1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106

use std::{pin::Pin, time::Instant};
use futures_lite::Future;

use tracing::{info, error, span::{Span}};
use tracing_futures::Instrument;

type Result<T> = anyhow::Result<T>;
pub type PinnedFut<'a, T=()> = Pin<Box<dyn Future<Output=Result<T>> + Send + 'a>>;

pub struct TracingTask<'a, R=()> {
    span: Span,
    future: PinnedFut<'a, R>,
    is_long_lived: bool
}

impl<'a, R: Send + Sync> TracingTask<'a, R> {
    pub fn new<T: Future<Output=Result<R>> + Send + 'a>(span: Span, fut: T) -> TracingTask<'a, R> {
        TracingTask {
            span,
            future: Box::pin(fut),
            is_long_lived: true
        }
    }

    pub fn new_short_lived<T: Future<Output=Result<R>> + Send + 'a>(span: Span, fut: T) -> TracingTask<'a, R> {
        TracingTask {
            span,
            future: Box::pin(fut),
            is_long_lived: false
        }
    }
}

impl<'a, R: Send + Sync + 'a> TracingTask<'a, R> {
    pub fn instrument(self) -> PinnedFut<'a, R> {
        let span = self.span;
        let future = self.future;
        let is_long_lived = self.is_long_lived;

        let fut_wrap = async move {
            if is_long_lived {
                info!("Starting...");
            }
            let t = Instant::now();

            let r = future.await;
            if r.is_err() {
                let err = r.err().unwrap();
                error!(error = ?err, elapsed = ?t.elapsed(), "Finished with");
                return Err(err);
            }
            info!(elapsed = ?t.elapsed(), "Finished [OK]...");
            Ok(r.unwrap())
        };

        Box::pin(fut_wrap.instrument(span))
    }
}

pub fn clean_fn(s: &str) -> String {
    let s = String::from(s);
    let name = s.split("::")
        .collect::<Vec<&str>>()
        .into_iter().rev()
        .take(2).rev()
        .collect::<Vec<&str>>()
        .join("::");

    let mut final_name = String::from("");
    let mut skip = 0;
    for c in name.chars() {
        if c == '<' {
            skip += 1;
        } else if c == '>' {
            skip -= 1;
        } else if skip < 1 {
            final_name.push(c);
        }
    }
    final_name
}

#[macro_export]
macro_rules! function {
    () => {{
        fn f() {}
        fn type_name_of<T>(_: T) -> &'static str {
            std::any::type_name::<T>()
        }
        let name = type_name_of(f);
        &name[..name.len() - 3]
    }}
}

#[macro_export]
macro_rules! span {
    ($($tts:tt)*) => {
        tracing::span!(tracing::Level::ERROR, "task", name = $crate::clean_fn($crate::function!()).as_str(), $($tts)*);
    };
    ($name:expr) => {
        tracing::span!(tracing::Level::ERROR, "task", name = $name);
    };
    () => {
        tracing::span!(tracing::Level::ERROR, "task", name = $crate::clean_fn($crate::function!()).as_str());
    };
}