use std::cell::Cell;
thread_local! {
static OPERATIONS: Cell<bool> = const { Cell::new(false) };
}
pub fn with_operation_spans<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
struct Restore(bool);
impl Drop for Restore {
fn drop(&mut self) {
OPERATIONS.with(|flag| flag.set(self.0));
}
}
let prev = OPERATIONS.with(|flag| flag.replace(enabled));
let _restore = Restore(prev);
f()
}
pub fn operations_enabled() -> bool {
OPERATIONS.with(Cell::get)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn off_by_default() {
assert!(!operations_enabled());
}
#[test]
fn on_inside_scope() {
with_operation_spans(true, || {
assert!(operations_enabled());
});
assert!(!operations_enabled());
}
#[test]
fn restores_after_panic() {
let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
with_operation_spans(true, || panic!("boom"));
}));
assert!(!operations_enabled());
}
}