1use pocket_ic::PocketIc;
2use std::{fmt, task::Poll, time::Duration};
3
4#[derive(Clone, Debug, Eq, PartialEq)]
6pub enum TickUntilError<E> {
7 ProgressLimit { rounds: u32 },
9 Failed(E),
11}
12
13impl<E: fmt::Display> fmt::Display for TickUntilError<E> {
14 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
15 match self {
16 Self::ProgressLimit { rounds } => {
17 write!(
18 formatter,
19 "readiness remained pending after {rounds} rounds"
20 )
21 }
22 Self::Failed(error) => write!(formatter, "readiness check failed: {error}"),
23 }
24 }
25}
26
27impl<E: std::error::Error + 'static> std::error::Error for TickUntilError<E> {
28 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
29 match self {
30 Self::ProgressLimit { .. } => None,
31 Self::Failed(error) => Some(error),
32 }
33 }
34}
35
36pub fn tick_until<T, E>(
48 pic: &PocketIc,
49 max_rounds: u32,
50 advance: Duration,
51 mut poll: impl FnMut(&PocketIc) -> Poll<Result<T, E>>,
52) -> Result<T, TickUntilError<E>> {
53 poll_with_rounds(
54 max_rounds,
55 || {
56 if !advance.is_zero() {
57 pic.advance_time(advance);
58 }
59 pic.tick();
60 },
61 || poll(pic),
62 )
63}
64
65fn poll_with_rounds<T, E>(
66 max_rounds: u32,
67 mut round: impl FnMut(),
68 mut poll: impl FnMut() -> Poll<Result<T, E>>,
69) -> Result<T, TickUntilError<E>> {
70 let mut rounds = 0;
71 loop {
72 match poll() {
73 Poll::Ready(result) => return result.map_err(TickUntilError::Failed),
74 Poll::Pending if rounds == max_rounds => {
75 return Err(TickUntilError::ProgressLimit { rounds });
76 }
77 Poll::Pending => {
78 round();
79 rounds += 1;
80 }
81 }
82 }
83}
84
85#[cfg(test)]
86mod tests {
87 use super::*;
88 use std::cell::Cell;
89
90 #[test]
91 fn initial_completion_and_failure_do_not_advance() {
92 for budget in [0, 3] {
93 for result in [Ok(42), Err("predicate failure")] {
94 assert_eq!(
95 poll_with_rounds(
96 budget,
97 || panic!("unexpected round"),
98 || Poll::Ready(result)
99 ),
100 result.map_err(TickUntilError::Failed)
101 );
102 }
103 }
104 }
105
106 #[test]
107 fn completion_on_last_round_preserves_poll_order() {
108 let rounds = Cell::new(0);
109 let polls = Cell::new(0);
110 let result = poll_with_rounds(
111 3,
112 || {
113 assert_eq!(polls.get(), rounds.get() + 1);
114 rounds.set(rounds.get() + 1);
115 },
116 || {
117 assert_eq!(polls.get(), rounds.get());
118 polls.set(polls.get() + 1);
119 if rounds.get() == 3 {
120 Poll::Ready(Ok::<_, ()>(42))
121 } else {
122 Poll::Pending
123 }
124 },
125 );
126 assert_eq!(result, Ok(42));
127 assert_eq!((rounds.get(), polls.get()), (3, 4));
128 }
129
130 #[test]
131 fn exhaustion_performs_exactly_the_allowed_rounds() {
132 for budget in [0, 1, 3] {
133 let rounds = Cell::new(0);
134 let polls = Cell::new(0);
135 let result = poll_with_rounds(
136 budget,
137 || rounds.set(rounds.get() + 1),
138 || {
139 polls.set(polls.get() + 1);
140 Poll::<Result<(), ()>>::Pending
141 },
142 );
143 assert_eq!(
144 result,
145 Err(TickUntilError::ProgressLimit { rounds: budget })
146 );
147 assert_eq!((rounds.get(), polls.get()), (budget, budget + 1));
148 }
149 }
150
151 #[test]
152 fn failure_after_progress_preserves_the_original_cause() {
153 let rounds = Cell::new(0);
154 let result = poll_with_rounds(
155 3,
156 || rounds.set(rounds.get() + 1),
157 || {
158 if rounds.get() == 1 {
159 Poll::Ready(Err::<(), _>(std::io::Error::from(
160 std::io::ErrorKind::PermissionDenied,
161 )))
162 } else {
163 Poll::Pending
164 }
165 },
166 );
167 let error = result.unwrap_err();
168 let source = std::error::Error::source(&error).unwrap();
169 assert_eq!(
170 source.downcast_ref::<std::io::Error>().unwrap().kind(),
171 std::io::ErrorKind::PermissionDenied
172 );
173 assert_eq!(rounds.get(), 1);
174 }
175 #[test]
176 fn polling_transport_skips_zero_time_and_preserves_nonzero_advancement() {
177 use std::{io::Write as _, net::TcpListener, thread, time::Instant};
178 for advance in [Duration::ZERO, Duration::from_nanos(7)] {
179 let listener = TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)).unwrap();
180 listener.set_nonblocking(true).unwrap();
181 let url = format!("http://{}/", listener.local_addr().unwrap())
182 .parse()
183 .unwrap();
184 let worker = thread::spawn(move || {
185 let deadline = Instant::now() + Duration::from_secs(5);
186 let mut requests = Vec::new();
187 loop {
188 assert!(Instant::now() < deadline, "polling peer timed out");
189 let (mut stream, _) = match listener.accept() {
190 Ok(connection) => connection,
191 Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
192 thread::sleep(Duration::from_millis(1));
193 continue;
194 }
195 Err(error) => panic!("accept polling request: {error}"),
196 };
197 let request = crate::testing::read_http_request(
198 &mut stream,
199 Duration::from_secs(1),
200 16_384,
201 16_384,
202 )
203 .unwrap();
204 let text = std::str::from_utf8(&request).unwrap();
205 let path = text.split_whitespace().nth(1).unwrap().to_owned();
206 let reply = if path.ends_with("/read/get_time") {
207 "{\"nanos_since_epoch\":100}"
208 } else {
209 "null"
210 };
211 if path.ends_with("/update/set_time") {
212 let body = text.split_once("\r\n\r\n").unwrap().1;
213 assert_eq!(
214 serde_json::from_str::<serde_json::Value>(body).unwrap()["nanos_since_epoch"],
215 100 + u64::from(advance.subsec_nanos())
216 );
217 }
218 stream
219 .set_write_timeout(Some(Duration::from_secs(1)))
220 .unwrap();
221 write!(stream, "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{reply}", reply.len()).unwrap();
222 let finished = path.ends_with("/update/tick");
223 requests.push(path);
224 if finished {
225 return requests;
226 }
227 }
228 });
229 let pic = PocketIc::new_from_existing_instance(url, 0, Some(1_000));
230 let mut polls = 0;
231 let result = tick_until(&pic, 1, advance, |_| {
232 polls += 1;
233 if polls == 2 {
234 Poll::Ready(Ok::<_, ()>(42))
235 } else {
236 Poll::Pending
237 }
238 });
239 assert_eq!(result, Ok(42));
240 assert_eq!(polls, 2);
241 let requests = worker.join().unwrap();
242 let expected = if advance.is_zero() {
243 vec!["/instances/0/update/tick"]
244 } else {
245 vec![
246 "/instances/0/read/get_time",
247 "/instances/0/update/set_time",
248 "/instances/0/update/tick",
249 ]
250 };
251 assert_eq!(requests, expected);
252 }
253 }
254}