1use std::collections::BTreeMap;
30use std::future::Future;
31use std::io::{BufRead, BufReader, Write};
32use std::net::{SocketAddr, TcpStream};
33use std::pin::Pin;
34use std::sync::atomic::{AtomicU64, Ordering};
35use std::sync::mpsc;
36use std::sync::{Arc, Mutex};
37use std::task::{Context, Poll};
38use std::time::Duration;
39
40use onetaskgraph_plugin_api::{Clock, SharedClock, system_clock};
41use tokio::sync::oneshot;
42
43use crate::Environment;
44
45pub use onetaskgraph_plugin_api::SIMULATED_CLOCK_VARIABLE;
48
49#[derive(Debug, Clone, PartialEq, Eq)]
51pub enum ClockChoice {
52 System,
54 Simulated {
56 address: SocketAddr,
58 client: usize,
60 },
61}
62
63pub fn clock_choice(value: Option<&str>) -> Result<ClockChoice, String> {
69 let value = value.unwrap_or_default().trim();
70 if value.is_empty() {
71 return Ok(ClockChoice::System);
72 }
73 let refuse = || {
74 format!(
75 "{SIMULATED_CLOCK_VARIABLE} is {value:?}, which is not `<address>/<client>` — a \
76 loopback address and a client number, as a test's simulated clock hands it out; \
77 unset it to run on the real clock"
78 )
79 };
80 let (address, client) = value.rsplit_once('/').ok_or_else(refuse)?;
81 let address: SocketAddr = address.parse().map_err(|_| refuse())?;
82 if !address.ip().is_loopback() {
85 return Err(format!(
86 "{SIMULATED_CLOCK_VARIABLE} names {address}, which is not a loopback address; a \
87 simulated clock is a test's own, on this machine — unset it to run on the real \
88 clock"
89 ));
90 }
91 Ok(ClockChoice::Simulated {
92 address,
93 client: client.parse().map_err(|_| refuse())?,
94 })
95}
96
97pub fn process_clock(environment: &Environment) -> Result<SharedClock, String> {
104 match clock_choice(environment.get(SIMULATED_CLOCK_VARIABLE))? {
105 ClockChoice::System => Ok(system_clock()),
106 ClockChoice::Simulated { address, client } => attach(address, client),
107 }
108}
109
110pub fn attach(address: SocketAddr, client: usize) -> Result<SharedClock, String> {
118 let unreachable = |error: std::io::Error| {
119 format!(
120 "{SIMULATED_CLOCK_VARIABLE} names a simulated clock at {address} that could not \
121 be reached: {error}; unset it to run on the real clock"
122 )
123 };
124 let stream = TcpStream::connect(address).map_err(unreachable)?;
125 stream.set_nodelay(true).map_err(unreachable)?;
126 let mut writer = stream.try_clone().map_err(unreachable)?;
127 let mut reader = BufReader::new(stream);
128 writeln!(writer, "attach {client}").map_err(unreachable)?;
129 writer.flush().map_err(unreachable)?;
130 let mut line = String::new();
131 reader.read_line(&mut line).map_err(unreachable)?;
132 if line.trim_end() != "attached" {
133 return Err(format!(
134 "the simulated clock at {address} answered {line:?} to an attach rather than \
135 `attached`; unset {SIMULATED_CLOCK_VARIABLE} to run on the real clock"
136 ));
137 }
138 let shared = Arc::new(Shared {
139 address,
140 writer: Mutex::new(writer),
141 pending: Mutex::new(Pending::default()),
142 sequence: AtomicU64::new(0),
143 });
144 let listening = Arc::clone(&shared);
145 std::thread::spawn(move || listening.listen(reader));
146 Ok(Arc::new(SimulatedClient { shared }))
147}
148
149struct SimulatedClient {
151 shared: Arc<Shared>,
152}
153
154struct Shared {
156 address: SocketAddr,
157 writer: Mutex<TcpStream>,
158 pending: Mutex<Pending>,
159 sequence: AtomicU64,
160}
161
162#[derive(Default)]
164struct Pending {
165 nows: BTreeMap<u64, mpsc::Sender<Duration>>,
166 wakes: BTreeMap<u64, oneshot::Sender<()>>,
167 closed: bool,
169 latest: Duration,
171}
172
173impl Shared {
174 fn send(&self, line: &str) {
177 let mut writer = self
178 .writer
179 .lock()
180 .expect("the clock's writer is not poisoned");
181 let _ = writeln!(writer, "{line}").and_then(|()| writer.flush());
182 }
183
184 fn lost(&self) -> ! {
188 panic!(
190 "the simulated clock at {} closed its connection or answered with something that \
191 is not a time, so this process's time is unknown; check the test that started it",
192 self.address
193 )
194 }
195
196 fn next(&self) -> u64 {
197 self.sequence.fetch_add(1, Ordering::Relaxed)
198 }
199
200 fn listen(&self, reader: BufReader<TcpStream>) {
202 for line in reader.lines() {
203 let Ok(line) = line else { break };
204 let mut words = line.split_whitespace();
205 let mut pending = self
206 .pending
207 .lock()
208 .expect("the clock's waiters are not poisoned");
209 let verb = words.next();
210 let seq = words.next().and_then(|seq| seq.parse::<u64>().ok());
211 match (verb, seq, words.next(), words.next()) {
212 (Some("now"), Some(seq), Some(nanos), None) => {
213 let Ok(nanos) = nanos.parse::<u64>() else {
216 break;
217 };
218 let now = Duration::from_nanos(nanos);
222 let Some(waiter) = pending.nows.remove(&seq) else {
223 break;
224 };
225 if now < pending.latest {
226 break;
227 }
228 pending.latest = now;
229 let _ = waiter.send(now);
230 }
231 (Some("wake"), Some(seq), None, None) => {
234 if let Some(waiter) = pending.wakes.remove(&seq) {
235 let _ = waiter.send(());
236 }
237 }
238 _ => break,
240 }
241 }
242 let mut pending = self
243 .pending
244 .lock()
245 .expect("the clock's waiters are not poisoned");
246 pending.closed = true;
247 pending.nows.clear();
248 pending.wakes.clear();
249 }
250}
251
252impl Clock for SimulatedClient {
253 fn now(&self) -> Duration {
254 let seq = self.shared.next();
255 let (sender, receiver) = mpsc::channel();
256 {
257 let mut pending = self
258 .shared
259 .pending
260 .lock()
261 .expect("the clock's waiters are not poisoned");
262 if pending.closed {
263 self.shared.lost();
264 }
265 pending.nows.insert(seq, sender);
266 }
267 self.shared.send(&format!("now {seq}"));
268 receiver.recv().unwrap_or_else(|_| self.shared.lost())
269 }
270
271 fn sleep(&self, duration: Duration) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
272 let seq = self.shared.next();
273 let (sender, receiver) = oneshot::channel();
274 {
275 let mut pending = self
276 .shared
277 .pending
278 .lock()
279 .expect("the clock's waiters are not poisoned");
280 if pending.closed {
281 self.shared.lost();
282 }
283 pending.wakes.insert(seq, sender);
284 }
285 self.shared
286 .send(&format!("sleep {seq} {}", duration.as_nanos()));
287 Box::pin(Wait {
288 shared: Arc::clone(&self.shared),
289 seq,
290 receiver,
291 woken: false,
292 })
293 }
294}
295
296struct Wait {
299 shared: Arc<Shared>,
300 seq: u64,
301 receiver: oneshot::Receiver<()>,
302 woken: bool,
303}
304
305impl Future for Wait {
306 type Output = ();
307
308 fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> {
309 match Pin::new(&mut self.receiver).poll(context) {
310 Poll::Ready(Ok(())) => {
311 self.woken = true;
312 Poll::Ready(())
313 }
314 Poll::Ready(Err(_)) => self.shared.lost(),
316 Poll::Pending => Poll::Pending,
317 }
318 }
319}
320
321impl Drop for Wait {
322 fn drop(&mut self) {
323 if !self.woken {
324 self.shared
325 .pending
326 .lock()
327 .expect("the clock's waiters are not poisoned")
328 .wakes
329 .remove(&self.seq);
330 self.shared.send(&format!("cancel {}", self.seq));
331 }
332 }
333}
334
335#[cfg(test)]
336mod tests {
337 use super::*;
338
339 #[test]
340 fn the_variable_unset_or_empty_selects_the_real_clock() {
341 assert_eq!(clock_choice(None), Ok(ClockChoice::System));
342 assert_eq!(clock_choice(Some("")), Ok(ClockChoice::System));
343 assert_eq!(clock_choice(Some(" ")), Ok(ClockChoice::System));
344 }
345
346 #[test]
347 fn a_value_names_the_coordinator_and_the_client() {
348 assert_eq!(
349 clock_choice(Some("127.0.0.1:4567/1")),
350 Ok(ClockChoice::Simulated {
351 address: "127.0.0.1:4567".parse().expect("an address"),
352 client: 1,
353 })
354 );
355 for malformed in [
356 "127.0.0.1:4567",
357 "nowhere/0",
358 "127.0.0.1:4567/x",
359 "192.0.2.1:4567/0",
360 ] {
361 let refused = clock_choice(Some(malformed)).expect_err("refused");
362 assert!(refused.contains(SIMULATED_CLOCK_VARIABLE), "{refused}");
363 }
364 }
365
366 #[tokio::test]
367 async fn an_unset_variable_builds_the_real_clock_which_sleeps_for_real() {
368 let clock = process_clock(&Environment::from_pairs(
369 std::iter::empty::<(String, String)>(),
370 ))
371 .expect("the real clock");
372 let before = clock.now();
373 clock.sleep(Duration::from_millis(20)).await;
374 assert!(clock.now() - before >= Duration::from_millis(20));
375 }
376}