Skip to main content

openinfer_simulator/
timer.rs

1use std::sync::{Mutex, OnceLock};
2use std::time::Instant;
3
4#[derive(Debug, Default)]
5struct TimerState {
6    starts: Vec<Option<Instant>>,
7    durations: Vec<u128>,
8    enabled: Vec<bool>,
9}
10
11/// Simple per-thread timing utility for runtime instrumentation.
12pub struct Timer;
13
14impl Timer {
15    fn state() -> &'static Mutex<TimerState> {
16        static INSTANCE: OnceLock<Mutex<TimerState>> = OnceLock::new();
17        INSTANCE.get_or_init(|| Mutex::new(TimerState::default()))
18    }
19
20    fn ensure_slot(state: &mut TimerState, thread_id: usize) {
21        if state.starts.len() <= thread_id {
22            state
23                .starts
24                .resize_with(thread_id + 1, || None);
25        }
26        if state.durations.len() <= thread_id {
27            state.durations.resize(thread_id + 1, 0);
28        }
29        if state.enabled.len() <= thread_id {
30            state.enabled.resize(thread_id + 1, false);
31        }
32    }
33
34    /// Enable or disable timing for a thread ID.
35    pub fn set_enabled(thread_id: usize, enabled: bool) {
36        let mut state = Self::state()
37            .lock()
38            .expect("timer state mutex poisoned");
39        Self::ensure_slot(&mut state, thread_id);
40        state.enabled[thread_id] = enabled;
41        if !enabled {
42            state.starts[thread_id] = None;
43            state.durations[thread_id] = 0;
44        }
45    }
46
47    /// Start a timer for a thread ID.
48    pub fn start(thread_id: usize) {
49        let mut state = Self::state()
50            .lock()
51            .expect("timer state mutex poisoned");
52        Self::ensure_slot(&mut state, thread_id);
53        if !state.enabled[thread_id] {
54            return;
55        }
56        state.starts[thread_id] = Some(Instant::now());
57    }
58
59    /// Stop a timer for a thread ID and record elapsed time.
60    pub fn stop(thread_id: usize) {
61        let mut state = Self::state()
62            .lock()
63            .expect("timer state mutex poisoned");
64        Self::ensure_slot(&mut state, thread_id);
65        if !state.enabled[thread_id] {
66            return;
67        }
68        let elapsed = state.starts[thread_id]
69            .take()
70            .map(|start| start.elapsed().as_nanos())
71            .unwrap_or(0);
72        state.durations[thread_id] = elapsed;
73    }
74
75    /// Record an explicit duration (ns) for a thread ID.
76    pub fn record(thread_id: usize, duration_ns: u128) {
77        let mut state = Self::state()
78            .lock()
79            .expect("timer state mutex poisoned");
80        Self::ensure_slot(&mut state, thread_id);
81        if !state.enabled[thread_id] {
82            return;
83        }
84        state.starts[thread_id] = None;
85        state.durations[thread_id] = duration_ns;
86    }
87
88    /// Fetch the last recorded duration (ns) for a thread ID.
89    pub fn elapsed(thread_id: usize) -> Option<u128> {
90        let state = Self::state()
91            .lock()
92            .expect("timer state mutex poisoned");
93        if thread_id >= state.enabled.len() || !state.enabled[thread_id] {
94            return Some(0);
95        }
96        Some(state.durations.get(thread_id).copied().unwrap_or(0))
97    }
98}