Skip to main content

muse2/
timeit.rs

1//! Implement a context manager for timing the decorated function
2use context_manager::{CallerContext, SyncWrapContext};
3use log::debug;
4use std::marker::PhantomData;
5use std::sync::{LazyLock, Mutex};
6use std::time::Instant;
7
8// Global variable to store the total time spent in investment step
9static INVESTMENT_TIME: LazyLock<Mutex<f64>> = LazyLock::new(|| Mutex::new(0.0));
10
11// Global variable to store the total time spent in dispatch steps
12static DISPATCH_TIME: LazyLock<Mutex<f64>> = LazyLock::new(|| Mutex::new(0.0));
13
14/// Trait providing the context of actions to be carried by the timer
15pub trait TimeContext {
16    /// Adds the label to be used in the logs
17    fn label(caller_context: &CallerContext) -> String;
18
19    /// Includes de logic to accumulate the time in a global variable
20    fn accumulate(milliseconds: &f64);
21
22    /// Returns the total time spent in this context
23    fn total_time() -> f64 {
24        0.0
25    }
26}
27
28/// Generic timing implementation shared by all timing contexts
29pub struct GenericTimer<C: TimeContext = GenericContext> {
30    start: Instant,
31    _context: PhantomData<C>, // See https://doc.rust-lang.org/nomicon/phantom-data.html
32}
33
34impl<T, C: TimeContext> SyncWrapContext<T> for GenericTimer<C> {
35    fn new() -> Self {
36        Self {
37            start: Instant::now(),
38            _context: PhantomData,
39        }
40    }
41
42    fn after(self, caller_context: &CallerContext, _result: &T) {
43        let label = C::label(caller_context);
44        let ms = self.start.elapsed().as_secs_f64() * 1000.0;
45        C::accumulate(&ms);
46        debug!("{label} completed in {ms:.1}ms");
47    }
48}
49
50/// Default context for timer, not accumularinging the time s
51/// uaing the decorated function name in the logs
52pub struct GenericContext;
53impl TimeContext for GenericContext {
54    fn label(caller_context: &CallerContext) -> String {
55        let fn_name = caller_context.fn_name();
56        format!("[Timer] '{fn_name}'")
57    }
58    fn accumulate(_milliseconds: &f64) {}
59}
60
61/// Investment context, using a fixed label in logs and accumulating
62/// the time in the `INVESTMENT_TIME` variable
63pub struct InvestmentContext;
64impl TimeContext for InvestmentContext {
65    fn label(_: &CallerContext) -> String {
66        "Investment step".to_string()
67    }
68    fn accumulate(milliseconds: &f64) {
69        let mut investment_time = INVESTMENT_TIME.lock().unwrap();
70        *investment_time += *milliseconds;
71    }
72    fn total_time() -> f64 {
73        let investment_time = INVESTMENT_TIME.lock().unwrap();
74        *investment_time
75    }
76}
77
78/// Dispatch context, using a fixed label in logs and accumulating
79/// the time in the `DISPATCH_TIME` variable
80pub struct DispatchContext;
81impl TimeContext for DispatchContext {
82    fn label(_: &CallerContext) -> String {
83        "Dispatch step".to_string()
84    }
85    fn accumulate(milliseconds: &f64) {
86        let mut dispatch_time = DISPATCH_TIME.lock().unwrap();
87        *dispatch_time += *milliseconds;
88    }
89    fn total_time() -> f64 {
90        let dispatch_time = DISPATCH_TIME.lock().unwrap();
91        *dispatch_time
92    }
93}
94
95/// Investment context for timing a function call
96pub type InvestmentTimer = GenericTimer<InvestmentContext>;
97
98/// Dispatch context for timing a function call
99pub type DispatchTimer = GenericTimer<DispatchContext>;