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