diff --git a/governor/src/state/in_memory.rs b/governor/src/state/in_memory.rs index 03f6802..269b52d 100644 --- a/governor/src/state/in_memory.rs +++ b/governor/src/state/in_memory.rs @@ -311,4 +311,40 @@ mod test { // Check if the state was updated as expected assert_eq!(state.0.load(Ordering::Acquire), 42); } + + #[cfg(feature = "std")] + #[test] + fn test_measure_and_peek_one_race_condition() { + let state = Arc::new(InMemoryState(AtomicU64::new(0))); + let barrier = Arc::new(std::sync::Barrier::new(2)); + + let state_clone_for_thread_a = Arc::clone(&state); + let barrier_clone_for_thread_a = Arc::clone(&barrier); + let thread_a = std::thread::spawn(move || { + barrier_clone_for_thread_a.wait(); // Ensure both threads start together + + state_clone_for_thread_a + .measure_and_peek_one(|old| { + // Simulate some work or delay if necessary + std::thread::sleep(Duration::from_millis(10)); + + let new_val = old.map(|n| n.as_u64() + 100).unwrap_or(100); + Ok::<(_, Nanos), &str>(((), Nanos::from(new_val))) + }) + .unwrap(); + }); + + let state_clone_for_thread_b = Arc::clone(&state); + let barrier_clone_for_thread_b = Arc::clone(&barrier); + let thread_b = std::thread::spawn(move || { + barrier_clone_for_thread_b.wait(); // Ensure both threads start together + + // Change the atomic variable's value after Thread A reads it + std::thread::sleep(Duration::from_millis(5)); + state_clone_for_thread_b.0.store(42, Ordering::Release); + }); + + thread_a.join().unwrap(); + thread_b.join().unwrap(); + } }