Skip to main content

tailtriage_core/
timers.rs

1use crate::collector::lock_map;
2use crate::{InFlightSnapshot, QueueEvent, StageEvent, Tailtriage};
3
4/// RAII guard tracking one in-flight unit for a named gauge.
5#[derive(Debug)]
6pub struct InflightGuard<'a> {
7    pub(crate) tailtriage: &'a Tailtriage,
8    pub(crate) gauge: String,
9    pub(crate) enabled: bool,
10}
11
12impl Drop for InflightGuard<'_> {
13    fn drop(&mut self) {
14        if !self.enabled {
15            return;
16        }
17
18        let count = {
19            let mut counts = lock_map(&self.tailtriage.inflight_counts);
20            let entry = counts.entry(self.gauge.clone()).or_insert(0);
21            if *entry > 0 {
22                *entry -= 1;
23            }
24            *entry
25        };
26
27        let sample = self.tailtriage.run_clock.sample();
28        self.tailtriage.record_inflight_snapshot(InFlightSnapshot {
29            gauge: self.gauge.clone(),
30            at_unix_ms: sample.unix_ms,
31            at_run_us: Some(sample.run_elapsed_us),
32            count,
33        });
34    }
35}
36
37/// Thin wrapper for recording stage latency around one await point.
38#[derive(Debug)]
39pub struct StageTimer<'a> {
40    pub(crate) tailtriage: &'a Tailtriage,
41    pub(crate) enabled: bool,
42    pub(crate) request_id: String,
43    pub(crate) stage: String,
44}
45
46impl StageTimer<'_> {
47    /// Awaits `fut`, records stage duration, and returns the original output.
48    ///
49    /// This helper is intended for fallible stage work where success can be
50    /// derived from `Result::is_ok`.
51    ///
52    /// Prefer this method when your stage naturally returns `Result<T, E>` and
53    /// you want success/failure evidence in the resulting triage report.
54    ///
55    /// # Errors
56    ///
57    /// Returns the same error value produced by `fut` after recording the
58    /// stage event with `success = false`.
59    pub async fn await_on<Fut, T, E>(self, fut: Fut) -> Result<T, E>
60    where
61        Fut: std::future::Future<Output = Result<T, E>>,
62    {
63        if !self.enabled {
64            return fut.await;
65        }
66
67        let interval_start = self.tailtriage.run_clock.start_interval();
68        let value = fut.await;
69        let finished = self.tailtriage.run_clock.finish_interval(interval_start);
70        let success = value.is_ok();
71
72        self.tailtriage.record_stage_event(StageEvent {
73            request_id: self.request_id,
74            stage: self.stage,
75            started_at_unix_ms: finished.started_at_unix_ms,
76            started_at_run_us: finished.started_at_run_us,
77            finished_at_unix_ms: finished.finished_at_unix_ms,
78            finished_at_run_us: finished.finished_at_run_us,
79            latency_us: finished.duration_us,
80            success,
81        });
82
83        value
84    }
85
86    /// Awaits an infallible stage future and records a successful stage event.
87    ///
88    /// Use this method when there is no meaningful stage-level error signal
89    /// (for example, internal CPU work or a prevalidated transformation).
90    pub async fn await_value<Fut, T>(self, fut: Fut) -> T
91    where
92        Fut: std::future::Future<Output = T>,
93    {
94        if !self.enabled {
95            return fut.await;
96        }
97
98        let interval_start = self.tailtriage.run_clock.start_interval();
99        let value = fut.await;
100        let finished = self.tailtriage.run_clock.finish_interval(interval_start);
101
102        self.tailtriage.record_stage_event(StageEvent {
103            request_id: self.request_id,
104            stage: self.stage,
105            started_at_unix_ms: finished.started_at_unix_ms,
106            started_at_run_us: finished.started_at_run_us,
107            finished_at_unix_ms: finished.finished_at_unix_ms,
108            finished_at_run_us: finished.finished_at_run_us,
109            latency_us: finished.duration_us,
110            success: true,
111        });
112
113        value
114    }
115}
116
117/// Thin wrapper for recording queue-wait latency around one await point.
118#[derive(Debug)]
119pub struct QueueTimer<'a> {
120    pub(crate) tailtriage: &'a Tailtriage,
121    pub(crate) enabled: bool,
122    pub(crate) request_id: String,
123    pub(crate) queue: String,
124    pub(crate) depth_at_start: Option<u64>,
125}
126
127impl QueueTimer<'_> {
128    /// Sets the queue depth sample captured at wait start.
129    #[must_use]
130    pub fn with_depth_at_start(mut self, depth_at_start: u64) -> Self {
131        self.depth_at_start = Some(depth_at_start);
132        self
133    }
134
135    /// Awaits `fut`, records queue wait duration, and returns the original output.
136    ///
137    /// Queue events are interpreted as application-level wait evidence (a lead,
138    /// not proof). Record these around bounded resources to help separate
139    /// queueing pressure from slow downstream stage time.
140    pub async fn await_on<Fut, T>(self, fut: Fut) -> T
141    where
142        Fut: std::future::Future<Output = T>,
143    {
144        if !self.enabled {
145            return fut.await;
146        }
147
148        let interval_start = self.tailtriage.run_clock.start_interval();
149        let value = fut.await;
150        let finished = self.tailtriage.run_clock.finish_interval(interval_start);
151
152        self.tailtriage.record_queue_event(QueueEvent {
153            request_id: self.request_id,
154            queue: self.queue,
155            waited_from_unix_ms: finished.started_at_unix_ms,
156            waited_from_run_us: finished.started_at_run_us,
157            waited_until_unix_ms: finished.finished_at_unix_ms,
158            waited_until_run_us: finished.finished_at_run_us,
159            wait_us: finished.duration_us,
160            depth_at_start: self.depth_at_start,
161        });
162
163        value
164    }
165}