darkfi/system/
timeout.rs

1/* This file is part of DarkFi (https://dark.fi)
2 *
3 * Copyright (C) 2020-2025 Dyne.org foundation
4 *
5 * This program is free software: you can redistribute it and/or modify
6 * it under the terms of the GNU Affero General Public License as
7 * published by the Free Software Foundation, either version 3 of the
8 * License, or (at your option) any later version.
9 *
10 * This program is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 * GNU Affero General Public License for more details.
14 *
15 * You should have received a copy of the GNU Affero General Public License
16 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
17 */
18
19use std::{
20    error::Error,
21    fmt,
22    future::Future,
23    io,
24    pin::Pin,
25    task::{Context, Poll},
26    time::Duration,
27};
28
29use pin_project_lite::pin_project;
30use smol::Timer;
31
32/// Awaits an I/O future or times out after a duration of time.
33///
34/// If you want to await a non I/O future consider using
35/// `timeout()` instead.
36///
37/// # Examples
38///
39/// ```no_run
40/// # fn main() -> std::io::Result<()> { smol::block_on(async {
41/// #
42/// use std::time::Duration;
43/// use std::io;
44///
45/// io_timeout(Duration::from_secs(5), async {
46///     let stdin = io::stdin();
47///     let mut line = String::new();
48///     let n = stdin.read_line(&mut line)?;
49///     Ok(())
50/// })
51/// .await?;
52/// #
53/// # Ok(()) }) }
54pub async fn io_timeout<F, T>(dur: Duration, f: F) -> io::Result<T>
55where
56    F: Future<Output = io::Result<T>>,
57{
58    Timeout { timeout: Timer::after(dur), future: f }.await
59}
60
61pin_project! {
62    #[derive(Debug)]
63    pub struct Timeout<F, T>
64    where
65        F: Future<Output = io::Result<T>>,
66    {
67        #[pin]
68        future: F,
69        #[pin]
70        timeout: Timer,
71    }
72}
73
74impl<F, T> Future for Timeout<F, T>
75where
76    F: Future<Output = io::Result<T>>,
77{
78    type Output = io::Result<T>;
79
80    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
81        let this = self.project();
82        match this.future.poll(cx) {
83            Poll::Pending => {}
84            other => return other,
85        }
86
87        if this.timeout.poll(cx).is_ready() {
88            let err = Err(io::Error::new(io::ErrorKind::TimedOut, "future timed out"));
89            Poll::Ready(err)
90        } else {
91            Poll::Pending
92        }
93    }
94}
95
96/// Awaits a future or times out after a duration of time.
97///
98/// If you want to await an I/O future consider using
99/// `io_timeout` instead.
100///
101/// # Examples
102///
103/// ```
104/// # fn main() -> std::io::Result<()> { smol::block_on(async {
105/// #
106/// use std::time::Duration;
107/// use smol::future;
108///
109/// let never = future::pending::<()>();
110/// let dur = Duration::from_millis(5);
111/// assert!(timeout(dur, never).await.is_err());
112/// #
113/// # Ok(()) }) }
114/// ```
115pub async fn timeout<F, T>(dur: Duration, f: F) -> Result<T, TimeoutError>
116where
117    F: Future<Output = T>,
118{
119    TimeoutFuture::new(f, dur).await
120}
121
122pin_project! {
123    /// A future that times out after a duration of time.
124    pub struct TimeoutFuture<F> {
125        #[pin]
126        future: F,
127        #[pin]
128        delay: Timer,
129    }
130}
131
132impl<F> TimeoutFuture<F> {
133    #[allow(dead_code)]
134    pub(super) fn new(future: F, dur: Duration) -> TimeoutFuture<F> {
135        TimeoutFuture { future, delay: Timer::after(dur) }
136    }
137}
138
139impl<F: Future> Future for TimeoutFuture<F> {
140    type Output = Result<F::Output, TimeoutError>;
141
142    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
143        let this = self.project();
144        match this.future.poll(cx) {
145            Poll::Ready(v) => Poll::Ready(Ok(v)),
146            Poll::Pending => match this.delay.poll(cx) {
147                Poll::Ready(_) => Poll::Ready(Err(TimeoutError { _private: () })),
148                Poll::Pending => Poll::Pending,
149            },
150        }
151    }
152}
153
154/// An error returned when a future times out.
155#[derive(Clone, Copy, Debug, Eq, PartialEq)]
156pub struct TimeoutError {
157    _private: (),
158}
159
160impl Error for TimeoutError {}
161
162impl fmt::Display for TimeoutError {
163    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
164        "future has timed out".fmt(f)
165    }
166}