Skip to main content

reifydb_runtime/sync/
waiter.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::{fmt, sync::Arc};
5
6use reifydb_value::value::duration::Duration;
7
8#[cfg(not(reifydb_single_threaded))]
9use crate::context::clock::{MockClock, TimerId};
10use crate::{
11	context::clock::{Clock, TimerWake},
12	sync::{condvar::Condvar, mutex::Mutex},
13};
14
15#[cfg(not(reifydb_single_threaded))]
16const MOCK_PARK_BACKSTOP: Duration = Duration::from_seconds_const(30);
17
18#[derive(Debug)]
19struct Signal {
20	notified: Mutex<bool>,
21	condvar: Condvar,
22}
23
24impl Signal {
25	fn new() -> Self {
26		Self {
27			notified: Mutex::new(false),
28			condvar: Condvar::new(),
29		}
30	}
31}
32
33impl TimerWake for Signal {
34	fn wake(&self) {
35		let _guard = self.notified.lock();
36		self.condvar.notify_all();
37	}
38}
39
40pub struct WaiterHandle {
41	signal: Arc<Signal>,
42	#[cfg_attr(reifydb_single_threaded, allow(dead_code))]
43	clock: Option<Clock>,
44	on_notify: Mutex<Option<Box<dyn FnOnce() + Send>>>,
45}
46
47impl fmt::Debug for WaiterHandle {
48	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
49		f.debug_struct("WaiterHandle").finish_non_exhaustive()
50	}
51}
52
53impl Default for WaiterHandle {
54	fn default() -> Self {
55		Self::new()
56	}
57}
58
59impl WaiterHandle {
60	pub fn new() -> Self {
61		Self {
62			signal: Arc::new(Signal::new()),
63			clock: None,
64			on_notify: Mutex::new(None),
65		}
66	}
67
68	pub fn on_clock(clock: Clock) -> Self {
69		Self {
70			signal: Arc::new(Signal::new()),
71			clock: Some(clock),
72			on_notify: Mutex::new(None),
73		}
74	}
75
76	pub fn with_callback(callback: Box<dyn FnOnce() + Send>) -> Self {
77		Self {
78			signal: Arc::new(Signal::new()),
79			clock: None,
80			on_notify: Mutex::new(Some(callback)),
81		}
82	}
83
84	pub fn notify(&self) {
85		let mut guard = self.signal.notified.lock();
86		*guard = true;
87		self.signal.condvar.notify_one();
88		drop(guard);
89		if let Some(callback) = self.on_notify.lock().take() {
90			callback();
91		}
92	}
93
94	pub fn wait_timeout(&self, timeout: Duration) -> bool {
95		#[cfg(not(reifydb_single_threaded))]
96		if let Some(mock) = self.clock.as_ref().and_then(Clock::as_mock) {
97			return self.wait_until_virtual_deadline(mock, timeout);
98		}
99
100		let mut guard = self.signal.notified.lock();
101		if *guard {
102			return true;
103		}
104		!self.signal.condvar.wait_for(&mut guard, timeout).timed_out()
105	}
106
107	#[cfg(not(reifydb_single_threaded))]
108	fn wait_until_virtual_deadline(&self, mock: &MockClock, timeout: Duration) -> bool {
109		let deadline = u64::try_from(mock.now().to_nanos())
110			.expect("mock clock is never before the Unix epoch")
111			.saturating_add(nanos_of(timeout));
112		let _timer = TimerGuard::register(mock, deadline, self.signal.clone());
113
114		let mut guard = self.signal.notified.lock();
115		loop {
116			if *guard {
117				return true;
118			}
119			if u64::try_from(mock.now().to_nanos()).expect("mock clock is never before the Unix epoch")
120				>= deadline
121			{
122				return false;
123			}
124			if self.signal.condvar.wait_for(&mut guard, MOCK_PARK_BACKSTOP).timed_out() {
125				panic!("mock clock never advanced past the park deadline");
126			}
127		}
128	}
129}
130
131#[cfg(not(reifydb_single_threaded))]
132fn nanos_of(timeout: Duration) -> u64 {
133	timeout.to_std().as_nanos().min(u64::MAX as u128) as u64
134}
135
136#[cfg(not(reifydb_single_threaded))]
137struct TimerGuard<'a> {
138	clock: &'a MockClock,
139	id: TimerId,
140}
141
142#[cfg(not(reifydb_single_threaded))]
143impl<'a> TimerGuard<'a> {
144	fn register(clock: &'a MockClock, deadline_nanos: u64, signal: Arc<Signal>) -> Self {
145		Self {
146			clock,
147			id: clock.register_timer(deadline_nanos, signal),
148		}
149	}
150}
151
152#[cfg(not(reifydb_single_threaded))]
153impl Drop for TimerGuard<'_> {
154	fn drop(&mut self) {
155		self.clock.cancel_timer(self.id);
156	}
157}
158
159#[cfg(test)]
160mod tests {
161	use std::sync::{
162		Arc,
163		atomic::{AtomicUsize, Ordering},
164	};
165
166	use super::*;
167
168	#[test]
169	fn callback_fires_exactly_once() {
170		let count = Arc::new(AtomicUsize::new(0));
171		let c = count.clone();
172		let waiter = WaiterHandle::with_callback(Box::new(move || {
173			c.fetch_add(1, Ordering::SeqCst);
174		}));
175
176		waiter.notify();
177		waiter.notify();
178
179		assert_eq!(count.load(Ordering::SeqCst), 1, "one-shot callback must fire exactly once");
180		assert!(
181			waiter.wait_timeout(Duration::from_milliseconds(0).unwrap()),
182			"an already-notified waiter returns immediately"
183		);
184	}
185}