Skip to main content

ic_testkit/pic/
time.rs

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