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 advance_and_tick(pic: &PocketIc, advance: Duration, ticks: u32) {
45 if !advance.is_zero() {
46 pic.advance_time(advance);
47 }
48 for _ in 0..ticks {
49 pic.tick();
50 }
51}
52
53pub fn tick_until<T, E>(
65 pic: &PocketIc,
66 max_rounds: u32,
67 advance: Duration,
68 mut poll: impl FnMut(&PocketIc) -> Poll<Result<T, E>>,
69) -> Result<T, TickUntilError<E>> {
70 poll_with_rounds(
71 max_rounds,
72 || advance_and_tick(pic, advance, 1),
73 || poll(pic),
74 )
75}
76
77fn poll_with_rounds<T, E>(
78 max_rounds: u32,
79 mut round: impl FnMut(),
80 mut poll: impl FnMut() -> Poll<Result<T, E>>,
81) -> Result<T, TickUntilError<E>> {
82 let mut rounds = 0;
83 loop {
84 match poll() {
85 Poll::Ready(result) => return result.map_err(TickUntilError::Failed),
86 Poll::Pending if rounds == max_rounds => {
87 return Err(TickUntilError::ProgressLimit { rounds });
88 }
89 Poll::Pending => {
90 round();
91 rounds += 1;
92 }
93 }
94 }
95}
96
97#[cfg(test)]
98mod tests {
99 use super::*;
100 use std::cell::Cell;
101
102 #[test]
103 fn initial_completion_and_failure_do_not_advance() {
104 for budget in [0, 3] {
105 for result in [Ok(42), Err("predicate failure")] {
106 assert_eq!(
107 poll_with_rounds(
108 budget,
109 || panic!("unexpected round"),
110 || Poll::Ready(result)
111 ),
112 result.map_err(TickUntilError::Failed)
113 );
114 }
115 }
116 }
117
118 #[test]
119 fn completion_on_last_round_preserves_poll_order() {
120 let rounds = Cell::new(0);
121 let polls = Cell::new(0);
122 let result = poll_with_rounds(
123 3,
124 || {
125 assert_eq!(polls.get(), rounds.get() + 1);
126 rounds.set(rounds.get() + 1);
127 },
128 || {
129 assert_eq!(polls.get(), rounds.get());
130 polls.set(polls.get() + 1);
131 if rounds.get() == 3 {
132 Poll::Ready(Ok::<_, ()>(42))
133 } else {
134 Poll::Pending
135 }
136 },
137 );
138 assert_eq!(result, Ok(42));
139 assert_eq!((rounds.get(), polls.get()), (3, 4));
140 }
141
142 #[test]
143 fn exhaustion_performs_exactly_the_allowed_rounds() {
144 for budget in [0, 1, 3] {
145 let rounds = Cell::new(0);
146 let polls = Cell::new(0);
147 let result = poll_with_rounds(
148 budget,
149 || rounds.set(rounds.get() + 1),
150 || {
151 polls.set(polls.get() + 1);
152 Poll::<Result<(), ()>>::Pending
153 },
154 );
155 assert_eq!(
156 result,
157 Err(TickUntilError::ProgressLimit { rounds: budget })
158 );
159 assert_eq!((rounds.get(), polls.get()), (budget, budget + 1));
160 }
161 }
162
163 #[test]
164 fn failure_after_progress_preserves_the_original_cause() {
165 let rounds = Cell::new(0);
166 let result = poll_with_rounds(
167 3,
168 || rounds.set(rounds.get() + 1),
169 || {
170 if rounds.get() == 1 {
171 Poll::Ready(Err::<(), _>(std::io::Error::from(
172 std::io::ErrorKind::PermissionDenied,
173 )))
174 } else {
175 Poll::Pending
176 }
177 },
178 );
179 let error = result.unwrap_err();
180 let source = std::error::Error::source(&error).unwrap();
181 assert_eq!(
182 source.downcast_ref::<std::io::Error>().unwrap().kind(),
183 std::io::ErrorKind::PermissionDenied
184 );
185 assert_eq!(rounds.get(), 1);
186 }
187 #[test]
188 fn polling_transport_skips_zero_time_and_preserves_nonzero_advancement() {
189 for advance in [Duration::ZERO, Duration::from_nanos(7)] {
190 assert_time_requests(advance, 1, |pic| {
191 let mut polls = 0;
192 let result = tick_until(pic, 1, advance, |_| {
193 polls += 1;
194 if polls == 2 {
195 Poll::Ready(Ok::<_, ()>(42))
196 } else {
197 Poll::Pending
198 }
199 });
200 assert_eq!(result, Ok(42));
201 assert_eq!(polls, 2);
202 });
203 }
204 }
205
206 #[test]
207 fn tick_batches_preserve_exact_transport_order_and_selected_counts() {
208 for advance in [Duration::ZERO, Duration::from_nanos(7)] {
209 for ticks in [0, 2, 4, 24] {
210 assert_time_requests(advance, ticks, |pic| advance_and_tick(pic, advance, ticks));
211 }
212 }
213 }
214
215 #[test]
216 fn install_cooldown_uses_exactly_two_ticks() {
217 use crate::pic::CanisterInstallExt as _;
218 for advance in [Duration::ZERO, Duration::from_nanos(7)] {
219 assert_time_requests(advance, 2, |pic| {
220 pic.wait_out_install_code_rate_limit(advance);
221 });
222 }
223 }
224
225 fn assert_time_requests(advance: Duration, ticks: u32, run: impl FnOnce(&PocketIc)) {
226 use std::{io::Write as _, net::TcpListener, thread, time::Instant};
227 let mut expected = Vec::new();
228 if !advance.is_zero() {
229 expected.extend(["/instances/0/read/get_time", "/instances/0/update/set_time"]);
230 }
231 expected.extend(std::iter::repeat_n(
232 "/instances/0/update/tick",
233 ticks as usize,
234 ));
235 let request_count = expected.len();
236 let listener = TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)).unwrap();
237 listener.set_nonblocking(true).unwrap();
238 let url = format!("http://{}/", listener.local_addr().unwrap())
239 .parse()
240 .unwrap();
241 let worker = thread::spawn(move || {
242 let deadline = Instant::now() + Duration::from_secs(5);
243 let mut requests = Vec::new();
244 while requests.len() < request_count {
245 assert!(Instant::now() < deadline, "simulated-time peer timed out");
246 let (mut stream, _) = match listener.accept() {
247 Ok(connection) => connection,
248 Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
249 thread::sleep(Duration::from_millis(1));
250 continue;
251 }
252 Err(error) => panic!("accept simulated-time request: {error}"),
253 };
254 let request = crate::testing::read_http_request(
255 &mut stream,
256 Duration::from_secs(1),
257 16_384,
258 16_384,
259 )
260 .unwrap();
261 let text = std::str::from_utf8(&request).unwrap();
262 let path = text.split_whitespace().nth(1).unwrap().to_owned();
263 let reply = if path.ends_with("/read/get_time") {
264 "{\"nanos_since_epoch\":100}"
265 } else {
266 "null"
267 };
268 if path.ends_with("/update/set_time") {
269 let body = text.split_once("\r\n\r\n").unwrap().1;
270 assert_eq!(
271 serde_json::from_str::<serde_json::Value>(body).unwrap()["nanos_since_epoch"],
272 100 + u64::from(advance.subsec_nanos())
273 );
274 }
275 stream
276 .set_write_timeout(Some(Duration::from_secs(1)))
277 .unwrap();
278 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();
279 requests.push(path);
280 }
281 requests
282 });
283 let pic = PocketIc::new_from_existing_instance(url, 0, Some(1_000));
284 run(&pic);
285 assert_eq!(worker.join().unwrap(), expected);
286 }
287}