Skip to main content

coven_foundation/
test_work.rs

1//! Work measured within one test's asynchronous and blocking calls.
2
3use std::{
4    cell::RefCell,
5    future::Future,
6    sync::{Arc, Mutex},
7};
8
9/// Measured work; unrelated tests never contribute to these counts.
10#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
11pub struct Counts {
12    /// Binary decoders entered.
13    pub decodes: usize,
14    /// Passphrase key derivations entered.
15    pub custody_derivations: usize,
16    /// Snapshot records and their evidence leaves decoded, before or during installation.
17    pub snapshot_decodes: usize,
18    /// Snapshot records and evidence leaves decoded while the application writer was held.
19    pub snapshot_decodes_with_writer: usize,
20    /// Byte ranges read from files on disk through `FileReader::read_at`.
21    pub file_range_reads: usize,
22    /// Bytes read from received objects kept in the input database.
23    pub retained_bytes_read: usize,
24    /// Bytes fed to object hashers.
25    pub object_bytes_hashed: usize,
26}
27
28/// The measurements shared by a task and its explicitly bound blocking calls.
29#[derive(Clone)]
30pub struct Counter(Arc<Mutex<Counts>>);
31
32#[derive(Clone)]
33struct Active {
34    counter: Counter,
35    writing: bool,
36}
37
38thread_local! {
39    // Clippy asks for a `const` initializer this already has where the
40    // target has no native thread locals, as Android: std builds those
41    // from a function, which the lint reads as the initializer.
42    #[cfg_attr(target_os = "android", allow(clippy::missing_const_for_thread_local))]
43    static ACTIVE: RefCell<Option<Active>> = const { RefCell::new(None) };
44}
45
46fn enter<R>(active: Option<Active>, run: impl FnOnce() -> R) -> R {
47    struct PreviousActive(Option<Active>);
48    impl Drop for PreviousActive {
49        fn drop(&mut self) {
50            ACTIVE.with(|active| *active.borrow_mut() = self.0.take());
51        }
52    }
53    let _restore = PreviousActive(ACTIVE.with(|current| current.replace(active)));
54    run()
55}
56
57impl Counter {
58    /// A counter with nothing measured yet, for one test.
59    pub fn new() -> Self {
60        Self(Arc::new(Mutex::new(Counts::default())))
61    }
62
63    /// Take the counts after all measured calls have finished.
64    pub fn take(&self) -> Counts {
65        std::mem::take(&mut *self.0.lock().expect("work counter poisoned"))
66    }
67
68    /// Measure a synchronous call, including calls it binds to another thread.
69    pub fn scope<R>(&self, run: impl FnOnce() -> R) -> R {
70        enter(Some(Active { counter: self.clone(), writing: false }), run)
71    }
72
73    /// Measure each poll without leaving a thread-local scope across an await.
74    pub async fn run<F: Future>(&self, future: F) -> F::Output {
75        let mut future = std::pin::pin!(future);
76        std::future::poll_fn(|context| self.scope(|| future.as_mut().poll(context))).await
77    }
78}
79
80/// Preserve the initiating measurement when dispatching a blocking call.
81pub fn bind<R>(run: impl FnOnce() -> R + Send + 'static) -> impl FnOnce() -> R + Send + 'static {
82    let active = ACTIVE.with(|active| active.borrow().clone());
83    move || enter(active, run)
84}
85
86/// Mark the actual application-writer boundary for the current measurement.
87pub fn writing<R>(run: impl FnOnce() -> R) -> R {
88    let active = ACTIVE.with(|active| active.borrow().clone()).map(|mut active| {
89        active.writing = true;
90        active
91    });
92    enter(active, run)
93}
94
95fn record(update: impl FnOnce(&mut Counts, bool)) {
96    ACTIVE.with(|active| {
97        if let Some(active) = &*active.borrow() {
98            update(&mut active.counter.0.lock().expect("work counter poisoned"), active.writing);
99        }
100    });
101}
102
103/// Record entry to a binary decoder.
104pub fn decode() {
105    record(|counts, _| counts.decodes += 1);
106}
107
108/// Record entry to an actual passphrase derivation.
109pub fn custody_derivation() {
110    record(|counts, _| counts.custody_derivations += 1);
111}
112
113/// Record decoding of snapshot records or their authenticated evidence leaves.
114pub fn snapshot_record_decode() {
115    record(|counts, writing| {
116        counts.snapshot_decodes += 1;
117        counts.snapshot_decodes_with_writer += usize::from(writing);
118    });
119}
120
121/// Record a byte range read from a file on disk.
122pub fn file_range_read() {
123    record(|counts, _| counts.file_range_reads += 1);
124}
125
126/// Record bytes read from a received object kept in the input database.
127pub fn retained_read(bytes: usize) {
128    record(|counts, _| counts.retained_bytes_read += bytes);
129}
130
131/// Record bytes fed to an object hasher.
132pub fn object_hashed(bytes: usize) {
133    record(|counts, _| counts.object_bytes_hashed += bytes);
134}
135
136#[cfg(test)]
137#[path = "test_work_tests.rs"]
138mod tests;