diff --git a/crates/common/src/capacity_scope.rs b/crates/common/src/capacity_scope.rs index 9f4ea49a2..5d24ffbbf 100644 --- a/crates/common/src/capacity_scope.rs +++ b/crates/common/src/capacity_scope.rs @@ -118,6 +118,7 @@ pub fn drain_global_dirty_scopes() -> Vec { mod tests { use super::*; use std::sync::Mutex; + use std::thread; fn test_lock() -> &'static Mutex<()> { static LOCK: OnceLock> = OnceLock::new(); @@ -127,14 +128,34 @@ mod tests { fn clear_capacity_scope_registry_for_test() { capacity_scope_registry() .lock() - .expect("capacity scope registry poisoned") + .unwrap_or_else(|poisoned| poisoned.into_inner()) .clear(); global_dirty_scope_registry() .lock() - .expect("global dirty scope registry poisoned") + .unwrap_or_else(|poisoned| poisoned.into_inner()) .clear(); } + fn poison_capacity_scope_registry_for_test() { + let _ = thread::spawn(|| { + let _guard = capacity_scope_registry() + .lock() + .expect("capacity scope registry lock should succeed"); + panic!("poison capacity scope registry"); + }) + .join(); + } + + fn poison_global_dirty_scope_registry_for_test() { + let _ = thread::spawn(|| { + let _guard = global_dirty_scope_registry() + .lock() + .expect("global dirty scope registry lock should succeed"); + panic!("poison global dirty scope registry"); + }) + .join(); + } + #[test] fn record_and_take_capacity_scope_round_trips() { let _guard = test_lock().lock().expect("test lock poisoned"); @@ -208,12 +229,33 @@ mod tests { ); } - let entries = capacity_scope_registry().lock().expect("capacity scope registry poisoned"); + let entries = capacity_scope_registry() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); assert!(entries.len() <= CAPACITY_SCOPE_REGISTRY_HARD_LIMIT); drop(entries); clear_capacity_scope_registry_for_test(); } + #[test] + fn record_capacity_scope_recovers_from_poisoned_registry() { + let _guard = test_lock().lock().expect("test lock poisoned"); + clear_capacity_scope_registry_for_test(); + poison_capacity_scope_registry_for_test(); + let token = Uuid::new_v4(); + let scope = CapacityScope { + disks: vec![CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }], + }; + + record_capacity_scope(token, scope.clone()); + + assert_eq!(take_capacity_scope(token), Some(scope)); + clear_capacity_scope_registry_for_test(); + } + #[test] fn record_and_drain_global_dirty_scope_round_trips() { let _guard = test_lock().lock().expect("test lock poisoned"); @@ -244,4 +286,23 @@ mod tests { assert!(drain_global_dirty_scopes().is_empty()); clear_capacity_scope_registry_for_test(); } + + #[test] + fn record_global_dirty_scope_recovers_from_poisoned_registry() { + let _guard = test_lock().lock().expect("test lock poisoned"); + clear_capacity_scope_registry_for_test(); + poison_global_dirty_scope_registry_for_test(); + + record_global_dirty_scope(CapacityScope { + disks: vec![CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }], + }); + + let drained = drain_global_dirty_scopes(); + assert_eq!(drained.len(), 1); + assert_eq!(drained[0].endpoint, "node-a"); + clear_capacity_scope_registry_for_test(); + } }