coven_foundation/
test_work.rs1use std::{
4 cell::RefCell,
5 future::Future,
6 sync::{Arc, Mutex},
7};
8
9#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
11pub struct Counts {
12 pub decodes: usize,
14 pub custody_derivations: usize,
16 pub snapshot_decodes: usize,
18 pub snapshot_decodes_with_writer: usize,
20 pub file_range_reads: usize,
22 pub retained_bytes_read: usize,
24 pub object_bytes_hashed: usize,
26}
27
28#[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 #[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 pub fn new() -> Self {
60 Self(Arc::new(Mutex::new(Counts::default())))
61 }
62
63 pub fn take(&self) -> Counts {
65 std::mem::take(&mut *self.0.lock().expect("work counter poisoned"))
66 }
67
68 pub fn scope<R>(&self, run: impl FnOnce() -> R) -> R {
70 enter(Some(Active { counter: self.clone(), writing: false }), run)
71 }
72
73 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
80pub 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
86pub 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
103pub fn decode() {
105 record(|counts, _| counts.decodes += 1);
106}
107
108pub fn custody_derivation() {
110 record(|counts, _| counts.custody_derivations += 1);
111}
112
113pub 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
121pub fn file_range_read() {
123 record(|counts, _| counts.file_range_reads += 1);
124}
125
126pub fn retained_read(bytes: usize) {
128 record(|counts, _| counts.retained_bytes_read += bytes);
129}
130
131pub 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;