omena_testkit/
instrumentation_session.rs

1//! Thread-local instrumentation counters for resolver and salsa query runs in tests.
2
3use std::cell::RefCell;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::{Arc, OnceLock};
6
7#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
8pub struct ResolverStyleIdentityIndexCountsV0 {
9    pub build_count: usize,
10    pub build_work_count: usize,
11}
12
13#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
14pub struct SalsaQueryRunCountsV0 {
15    pub digest: usize,
16    pub dependency: usize,
17    pub transitive_leaf: usize,
18    pub transitive_a: usize,
19    pub transitive_b: usize,
20    pub transitive_c: usize,
21    pub transitive_unrelated: usize,
22}
23
24#[derive(Debug, Clone, Default)]
25pub struct InstrumentationSessionV0 {
26    inner: Arc<InstrumentationSessionInnerV0>,
27}
28
29#[derive(Debug, Default)]
30struct InstrumentationSessionInnerV0 {
31    resolver_identity_index_build_count: AtomicUsize,
32    resolver_identity_index_build_work_count: AtomicUsize,
33    salsa_digest_query_runs: AtomicUsize,
34    salsa_dependency_query_runs: AtomicUsize,
35    salsa_transitive_leaf_query_runs: AtomicUsize,
36    salsa_transitive_a_query_runs: AtomicUsize,
37    salsa_transitive_b_query_runs: AtomicUsize,
38    salsa_transitive_c_query_runs: AtomicUsize,
39    salsa_transitive_unrelated_query_runs: AtomicUsize,
40}
41
42thread_local! {
43    static CURRENT_INSTRUMENTATION_SESSION: RefCell<Option<InstrumentationSessionV0>> =
44        const { RefCell::new(None) };
45}
46
47static DEFAULT_INSTRUMENTATION_SESSION: OnceLock<InstrumentationSessionV0> = OnceLock::new();
48
49pub fn default_instrumentation_session_v0() -> InstrumentationSessionV0 {
50    DEFAULT_INSTRUMENTATION_SESSION
51        .get_or_init(InstrumentationSessionV0::default)
52        .clone()
53}
54
55pub fn current_instrumentation_session_v0() -> InstrumentationSessionV0 {
56    CURRENT_INSTRUMENTATION_SESSION
57        .with(|slot| slot.borrow().clone())
58        .unwrap_or_else(default_instrumentation_session_v0)
59}
60
61pub fn with_instrumentation_session<R>(
62    session: InstrumentationSessionV0,
63    body: impl FnOnce() -> R,
64) -> R {
65    let guard = InstalledInstrumentationSessionV0::new(session);
66    let result = body();
67    drop(guard);
68    result
69}
70
71struct InstalledInstrumentationSessionV0 {
72    previous: Option<InstrumentationSessionV0>,
73}
74
75impl InstalledInstrumentationSessionV0 {
76    fn new(session: InstrumentationSessionV0) -> Self {
77        let previous = CURRENT_INSTRUMENTATION_SESSION.with(|slot| slot.replace(Some(session)));
78        Self { previous }
79    }
80}
81
82impl Drop for InstalledInstrumentationSessionV0 {
83    fn drop(&mut self) {
84        let previous = self.previous.take();
85        CURRENT_INSTRUMENTATION_SESSION.with(|slot| {
86            let _ = slot.replace(previous);
87        });
88    }
89}
90
91impl InstrumentationSessionV0 {
92    pub fn reset_resolver_style_identity_index_counts(&self) {
93        self.inner
94            .resolver_identity_index_build_count
95            .store(0, Ordering::Release);
96        self.inner
97            .resolver_identity_index_build_work_count
98            .store(0, Ordering::Release);
99    }
100
101    pub fn record_resolver_style_identity_index_build(&self, work_count: usize) {
102        self.inner
103            .resolver_identity_index_build_count
104            .fetch_add(1, Ordering::AcqRel);
105        self.inner
106            .resolver_identity_index_build_work_count
107            .fetch_add(work_count, Ordering::AcqRel);
108    }
109
110    pub fn resolver_style_identity_index_counts(&self) -> ResolverStyleIdentityIndexCountsV0 {
111        ResolverStyleIdentityIndexCountsV0 {
112            build_count: self
113                .inner
114                .resolver_identity_index_build_count
115                .load(Ordering::Acquire),
116            build_work_count: self
117                .inner
118                .resolver_identity_index_build_work_count
119                .load(Ordering::Acquire),
120        }
121    }
122
123    pub fn reset_salsa_query_run_counts(&self) {
124        self.inner
125            .salsa_digest_query_runs
126            .store(0, Ordering::Release);
127        self.inner
128            .salsa_dependency_query_runs
129            .store(0, Ordering::Release);
130        self.inner
131            .salsa_transitive_leaf_query_runs
132            .store(0, Ordering::Release);
133        self.inner
134            .salsa_transitive_a_query_runs
135            .store(0, Ordering::Release);
136        self.inner
137            .salsa_transitive_b_query_runs
138            .store(0, Ordering::Release);
139        self.inner
140            .salsa_transitive_c_query_runs
141            .store(0, Ordering::Release);
142        self.inner
143            .salsa_transitive_unrelated_query_runs
144            .store(0, Ordering::Release);
145    }
146
147    pub fn record_salsa_digest_query_run(&self) {
148        self.inner
149            .salsa_digest_query_runs
150            .fetch_add(1, Ordering::AcqRel);
151    }
152
153    pub fn record_salsa_dependency_query_run(&self) {
154        self.inner
155            .salsa_dependency_query_runs
156            .fetch_add(1, Ordering::AcqRel);
157    }
158
159    pub fn record_salsa_transitive_leaf_query_run(&self) {
160        self.inner
161            .salsa_transitive_leaf_query_runs
162            .fetch_add(1, Ordering::AcqRel);
163    }
164
165    pub fn record_salsa_transitive_a_query_run(&self) {
166        self.inner
167            .salsa_transitive_a_query_runs
168            .fetch_add(1, Ordering::AcqRel);
169    }
170
171    pub fn record_salsa_transitive_b_query_run(&self) {
172        self.inner
173            .salsa_transitive_b_query_runs
174            .fetch_add(1, Ordering::AcqRel);
175    }
176
177    pub fn record_salsa_transitive_c_query_run(&self) {
178        self.inner
179            .salsa_transitive_c_query_runs
180            .fetch_add(1, Ordering::AcqRel);
181    }
182
183    pub fn record_salsa_transitive_unrelated_query_run(&self) {
184        self.inner
185            .salsa_transitive_unrelated_query_runs
186            .fetch_add(1, Ordering::AcqRel);
187    }
188
189    pub fn salsa_query_run_counts(&self) -> SalsaQueryRunCountsV0 {
190        SalsaQueryRunCountsV0 {
191            digest: self.inner.salsa_digest_query_runs.load(Ordering::Acquire),
192            dependency: self
193                .inner
194                .salsa_dependency_query_runs
195                .load(Ordering::Acquire),
196            transitive_leaf: self
197                .inner
198                .salsa_transitive_leaf_query_runs
199                .load(Ordering::Acquire),
200            transitive_a: self
201                .inner
202                .salsa_transitive_a_query_runs
203                .load(Ordering::Acquire),
204            transitive_b: self
205                .inner
206                .salsa_transitive_b_query_runs
207                .load(Ordering::Acquire),
208            transitive_c: self
209                .inner
210                .salsa_transitive_c_query_runs
211                .load(Ordering::Acquire),
212            transitive_unrelated: self
213                .inner
214                .salsa_transitive_unrelated_query_runs
215                .load(Ordering::Acquire),
216        }
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223    use std::sync::{Arc, Barrier};
224    use std::thread;
225
226    #[test]
227    fn instrumentation_sessions_isolate_concurrent_counts() -> Result<(), Box<dyn std::error::Error>>
228    {
229        let first = InstrumentationSessionV0::default();
230        let second = InstrumentationSessionV0::default();
231        let barrier = Arc::new(Barrier::new(2));
232
233        let first_worker = {
234            let session = first.clone();
235            let barrier = Arc::clone(&barrier);
236            thread::spawn(move || {
237                with_instrumentation_session(session, || {
238                    barrier.wait();
239                    for _ in 0..11 {
240                        current_instrumentation_session_v0()
241                            .record_resolver_style_identity_index_build(3);
242                    }
243                });
244            })
245        };
246        let second_worker = {
247            let session = second.clone();
248            let barrier = Arc::clone(&barrier);
249            thread::spawn(move || {
250                with_instrumentation_session(session, || {
251                    barrier.wait();
252                    for _ in 0..7 {
253                        current_instrumentation_session_v0()
254                            .record_resolver_style_identity_index_build(5);
255                    }
256                });
257            })
258        };
259
260        first_worker
261            .join()
262            .map_err(|_| std::io::Error::other("first instrumentation worker panicked"))?;
263        second_worker
264            .join()
265            .map_err(|_| std::io::Error::other("second instrumentation worker panicked"))?;
266
267        assert_eq!(
268            first.resolver_style_identity_index_counts(),
269            ResolverStyleIdentityIndexCountsV0 {
270                build_count: 11,
271                build_work_count: 33
272            }
273        );
274        assert_eq!(
275            second.resolver_style_identity_index_counts(),
276            ResolverStyleIdentityIndexCountsV0 {
277                build_count: 7,
278                build_work_count: 35
279            }
280        );
281        Ok(())
282    }
283
284    #[test]
285    fn instrumentation_session_restores_previous_session() {
286        let outer = InstrumentationSessionV0::default();
287        let inner = InstrumentationSessionV0::default();
288
289        with_instrumentation_session(outer.clone(), || {
290            current_instrumentation_session_v0().record_salsa_digest_query_run();
291            with_instrumentation_session(inner.clone(), || {
292                current_instrumentation_session_v0().record_salsa_digest_query_run();
293                current_instrumentation_session_v0().record_salsa_dependency_query_run();
294            });
295            current_instrumentation_session_v0().record_salsa_dependency_query_run();
296        });
297
298        assert_eq!(
299            outer.salsa_query_run_counts(),
300            SalsaQueryRunCountsV0 {
301                digest: 1,
302                dependency: 1,
303                ..SalsaQueryRunCountsV0::default()
304            }
305        );
306        assert_eq!(
307            inner.salsa_query_run_counts(),
308            SalsaQueryRunCountsV0 {
309                digest: 1,
310                dependency: 1,
311                ..SalsaQueryRunCountsV0::default()
312            }
313        );
314    }
315}
omena_testkit/instrumentation_session.rs

omena_testkit/
instrumentation_session.rs
