kairoi/
span.rs

1use crate::{Event, Node};
2use lazy_static::lazy_static;
3use std::sync::{Arc, Weak};
4use std::time::SystemTime;
5use tokio::task_local;
6
7#[derive(Debug, Copy, Clone)]
8pub struct Progress {
9    total: u64,
10    progress: u64,
11}
12
13impl Progress {
14    pub fn new(total: u64, progress: u64) -> Self {
15        Self { total, progress }
16    }
17
18    pub fn total(&self) -> u64 {
19        self.total
20    }
21
22    pub fn progress(&self) -> u64 {
23        self.progress
24    }
25}
26
27#[derive(Debug, Clone)]
28pub struct Span {
29    timestamp: SystemTime,
30    name: String,
31    progress: Option<Progress>,
32}
33
34pub type SpanRef = Arc<Node<Span>>;
35
36impl Span {
37    pub fn with_name(&self, name: String) -> Self {
38        let mut clone = self.clone();
39        clone.name = name;
40        clone
41    }
42
43    pub fn with_progress(&self, progress: Progress) -> Self {
44        let mut clone = self.clone();
45        clone.progress = Some(progress);
46        clone
47    }
48
49    pub fn timestamp(&self) -> SystemTime {
50        self.timestamp
51    }
52
53    pub fn name(&self) -> &str {
54        &self.name
55    }
56
57    pub fn progress(&self) -> Option<Progress> {
58        self.progress
59    }
60}
61
62impl Default for Span {
63    fn default() -> Self {
64        Self {
65            timestamp: SystemTime::now(),
66            name: String::default(),
67            progress: None,
68        }
69    }
70}
71
72lazy_static! {
73    static ref ROOT: SpanRef = Arc::new(Node::new(Span::default()));
74}
75
76task_local! {
77    static CURRENT: Weak<Node<Span >>;
78}
79
80impl Span {
81    pub async fn scope<T, F: AsyncFnOnce(Scope) -> T>(f: F) -> T {
82        let new = Self::current().add(Node::new(Self::default()));
83
84        Event::span_begin(new.clone()).submit();
85
86        let new_clone = new.clone();
87        let v = CURRENT
88            .scope(Arc::downgrade(&new_clone), async move {
89                f(Scope::new(new_clone)).await
90            })
91            .await;
92
93        new.delete();
94        Event::span_end(new).submit();
95
96        v
97    }
98
99    pub fn current() -> SpanRef {
100        CURRENT
101            .try_with(|v| v.upgrade())
102            .ok()
103            .flatten()
104            .unwrap_or_else(|| ROOT.clone())
105    }
106
107    pub fn root() -> SpanRef {
108        ROOT.clone()
109    }
110}
111
112pub struct Scope {
113    node: SpanRef,
114}
115
116impl Scope {
117    fn new(node: SpanRef) -> Self {
118        Self { node }
119    }
120
121    pub fn update(&self, data: Span) {
122        self.node.update(data);
123    }
124}
kairoi/span.rs

kairoi/
span.rs
