latency_trace/
probed_trace.rs

1//! Provides the ability to obtain interim timing information before the target function terminates.
2
3use crate::{LatencyTrace, Timings};
4use std::{
5    sync::{Arc, Mutex},
6    thread::JoinHandle,
7};
8
9/// Represents an ongoing collection of latency information with the ability to report on partial latencies
10/// before the instrumented function completes.
11#[derive(Clone)]
12pub struct ProbedTrace {
13    lt: LatencyTrace,
14    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
15}
16
17impl ProbedTrace {
18    pub(crate) fn new(lt: LatencyTrace) -> Self {
19        Self {
20            lt,
21            join_handle: Mutex::new(None).into(),
22        }
23    }
24
25    pub(crate) fn set_join_handle(&self, join_handle: JoinHandle<()>) {
26        let mut lock = self.join_handle.lock();
27        let jh = lock
28            .as_deref_mut()
29            .expect("ProbedTrace join_handle Mutex poisoned");
30        *jh = Some(join_handle);
31    }
32
33    /// Returns partial latencies collected when the call is made.
34    pub fn probe_latencies(&self) -> Timings {
35        let acc = self.lt.0.control.probe_tls();
36        self.lt.0.report_timings(acc)
37    }
38
39    /// Blocks until the function being measured completes, and then returns the collected latency information.
40    ///
41    /// Should only be called at most once, from main thread. May panic otherwise.
42    pub fn wait_and_report(&self) -> Timings {
43        // try_lock() below should always succeed because this function is the only one that should be joining
44        // the handle and it should only be called once from the main thread.
45        let join_handle = self
46            .join_handle
47            .try_lock()
48            .expect("ProbedTrace lock should not be contended")
49            .take()
50            .expect("`join_handle` set by constructor, may only be taken once");
51        join_handle
52            .join()
53            .expect("ProbedTrace execution thread exited abnormally");
54        let acc = self.lt.0.take_acc_timings();
55        self.lt.0.report_timings(acc)
56    }
57}