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/// Check a caller-owned predicate, advancing simulated time and ticking while pending.
37///
38/// The predicate runs once before any mutation. Each pending result permits one
39/// round of `advance_time(advance)`, then `tick()`, then another predicate check,
40/// up to `max_rounds`. A zero budget still performs the initial check. A zero
41/// duration permits ticks without reading or writing simulated time. Completion
42/// on the last allowed round succeeds; predicate failures return immediately without another round.
43///
44/// This bounds simulated progression, not wall-clock time spent inside PocketIC
45/// calls or the predicate. PocketIC operation failures retain upstream panic
46/// behavior. Applications own readiness meaning and predicate error types.
47pub 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}