Files
crank/crates/crank-artifacts/src/test_support.rs
T

128 lines
3.6 KiB
Rust

//! Test-only crash and fault control for syscall-coupled storage checkpoints.
use std::{
collections::BTreeMap,
sync::{
Condvar, Mutex, OnceLock,
atomic::{AtomicUsize, Ordering},
},
time::Duration,
};
static TRAVERSAL_CALLS: AtomicUsize = AtomicUsize::new(0);
pub fn reset_traversal_calls() {
TRAVERSAL_CALLS.store(0, Ordering::SeqCst);
}
pub fn traversal_calls() -> usize {
TRAVERSAL_CALLS.load(Ordering::SeqCst)
}
pub(crate) fn record_traversal_call() {
TRAVERSAL_CALLS.fetch_add(1, Ordering::SeqCst);
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum FaultAction {
Exit,
Fail,
Hold,
}
struct State {
checkpoint: Option<(String, FaultAction, usize)>,
hits: BTreeMap<String, usize>,
held: bool,
}
fn slot() -> &'static (Mutex<State>, Condvar) {
static SLOT: OnceLock<(Mutex<State>, Condvar)> = OnceLock::new();
SLOT.get_or_init(|| {
(
Mutex::new(State {
checkpoint: None,
hits: BTreeMap::new(),
held: false,
}),
Condvar::new(),
)
})
}
pub fn set_checkpoint(stage: impl Into<String>, action: FaultAction) {
set_checkpoint_on_hit(stage, action, 1);
}
pub fn set_checkpoint_on_hit(stage: impl Into<String>, action: FaultAction, hit: usize) {
assert!(hit > 0, "checkpoint hit is one-based");
let (lock, _) = slot();
let mut state = lock.lock().expect("artifact test checkpoint lock poisoned");
state.checkpoint = Some((stage.into(), action, hit));
state.hits.clear();
}
pub fn checkpoint_hits(stage: &str) -> usize {
let (lock, _) = slot();
lock.lock()
.expect("artifact test checkpoint lock poisoned")
.hits
.get(stage)
.copied()
.unwrap_or(0)
}
pub fn clear_checkpoint() {
let (lock, wake) = slot();
let mut state = lock.lock().expect("artifact test checkpoint lock poisoned");
state.checkpoint = None;
state.held = false;
wake.notify_all();
}
pub fn wait_until_held(timeout: Duration) -> bool {
let (lock, wake) = slot();
let state = lock.lock().expect("artifact test checkpoint lock poisoned");
let (state, _) = wake
.wait_timeout_while(state, timeout, |state| !state.held)
.expect("artifact test checkpoint lock poisoned");
state.held
}
pub(crate) fn checkpoint(stage: &str) -> Option<FaultAction> {
let (lock, wake) = slot();
let mut state = lock.lock().expect("artifact test checkpoint lock poisoned");
let configured = state.checkpoint.clone();
let action = configured.and_then(|(expected, action, target_hit)| {
if expected != stage {
return None;
}
let hit = state.hits.entry(stage.to_owned()).or_insert(0);
if *hit == 0 {
*hit = 1;
}
(*hit == target_hit).then_some(action)
});
if action == Some(FaultAction::Hold) {
state.held = true;
wake.notify_all();
while state.checkpoint.is_some() {
state = wake
.wait(state)
.expect("artifact test checkpoint lock poisoned");
}
}
action
}
pub(crate) fn action(stage: &str) -> Option<FaultAction> {
let (lock, _) = slot();
let mut state = lock.lock().expect("artifact test checkpoint lock poisoned");
*state.hits.entry(stage.to_owned()).or_insert(0) += 1;
let hit = state.hits[stage];
state
.checkpoint
.as_ref()
.and_then(|(expected, action, target_hit)| {
(expected == stage && hit == *target_hit).then_some(*action)
})
}