Skip to main content

ic_testkit/pic/
time.rs

1use pocket_ic::PocketIc;
2use std::{task::Poll, time::Duration};
3
4/// A caller-owned readiness check failed or exhausted its simulated round budget.
5#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
6pub enum TickUntilError<E> {
7    /// The predicate was still pending after this many advance/tick rounds.
8    #[error("readiness remained pending after {rounds} rounds")]
9    ProgressLimit { rounds: u32 },
10    /// The predicate returned its own failure.
11    #[error("readiness check failed: {0}")]
12    Failed(#[source] E),
13}
14
15/// Advance simulated time once, then deliver exactly the caller-selected ticks.
16///
17/// Zero duration performs no time read or write; zero ticks permits advancement
18/// without delivering a round. This function does not poll, retry, start a server
19/// or decide application readiness. Callers own completion and cadence budgets.
20///
21/// # Panics
22/// PocketIC operation failures retain the upstream client's panic behavior.
23pub 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
32/// Check a caller-owned predicate, advancing simulated time and ticking while pending.
33///
34/// The predicate runs once before any mutation. Each pending result permits one
35/// round of `advance_time(advance)`, then `tick()`, then another predicate check,
36/// up to `max_rounds`. A zero budget still performs the initial check. A zero
37/// duration permits ticks without reading or writing simulated time. Completion
38/// on the last allowed round succeeds; predicate failures return immediately without another round.
39///
40/// This bounds simulated progression, not wall-clock time spent inside PocketIC
41/// calls or the predicate. PocketIC operation failures retain upstream panic
42/// behavior. Applications own readiness meaning and predicate error types.
43pub 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}