smol_timeout/
lib.rs

1/* ┌────────────────────────────────────────────────────────────────────────────────────────────┐ *\
2 * │                                                                                            │ *
3 * │ This Source Code Form is subject to the terms of the Mozilla Public                        │ *
4 * │ License, v. 2.0. If a copy of the MPL was not distributed with this                        │ *
5 * │ file, You can obtain one at http://mozilla.org/MPL/2.0/.                                   │ *
6 * │                                                                                            │ *
7\* └────────────────────────────────────────────────────────────────────────────────────────────┘ */
8
9/* ┌────────────────────────────────────────────────────────────────────────────────────────────┐ *\
10 * │                                       Configuration                                        │ *
11\* └────────────────────────────────────────────────────────────────────────────────────────────┘ */
12
13#![no_std]
14
15/* ┌────────────────────────────────────────────────────────────────────────────────────────────┐ *\
16 * │                                       Documentation                                        │ *
17\* └────────────────────────────────────────────────────────────────────────────────────────────┘ */
18
19//! A way to poll a future until it or a timer completes.
20//!
21//! ## Example
22//!
23//! ```rust
24//! use async_io::Timer;
25//! # use futures_lite::future;
26//! use smol_timeout::TimeoutExt;
27//! use std::time::Duration;
28//!
29//! # future::block_on(async {
30//! #
31//! let foo = async {
32//!     Timer::after(Duration::from_millis(250)).await;
33//!     24
34//! };
35//!
36//! let foo = foo.timeout(Duration::from_millis(100));
37//! assert_eq!(foo.await, None);
38//!
39//! let bar = async {
40//!     Timer::after(Duration::from_millis(100)).await;
41//!     42
42//! };
43//!
44//! let bar = bar.timeout(Duration::from_millis(250));
45//! assert_eq!(bar.await, Some(42));
46//! #
47//! # });
48//! ```
49
50/* ┌────────────────────────────────────────────────────────────────────────────────────────────┐ *\
51 * │                                          Imports                                           │ *
52\* └────────────────────────────────────────────────────────────────────────────────────────────┘ */
53
54use async_io::Timer;
55use core::future::Future;
56use core::pin::Pin;
57use core::task::{Context, Poll};
58use core::time::Duration;
59use pin_project_lite::pin_project;
60
61/* ┌────────────────────────────────────────────────────────────────────────────────────────────┐ *\
62 * │                                    struct Timeout<Fut>                                     │ *
63\* └────────────────────────────────────────────────────────────────────────────────────────────┘ */
64
65pin_project! {
66    #[derive(Debug)]
67    /// A future polling both another future and a [`Timer`] that will complete after a specified
68    /// timeout, and returning the future's output or [`None`] if the timer completes first.
69    ///
70    /// ## Example
71    ///
72    /// ```rust
73    /// use async_io::Timer;
74    /// # use futures_lite::future;
75    /// use smol_timeout::TimeoutExt;
76    /// use std::time::Duration;
77    ///
78    /// # future::block_on(async {
79    /// #
80    /// let foo = async {
81    ///     Timer::after(Duration::from_millis(250)).await;
82    ///     24
83    /// };
84    ///
85    /// let foo = foo.timeout(Duration::from_millis(100));
86    /// assert_eq!(foo.await, None);
87    ///
88    /// let bar = async {
89    ///     Timer::after(Duration::from_millis(100)).await;
90    ///     42
91    /// };
92    ///
93    /// let bar = bar.timeout(Duration::from_millis(250));
94    /// assert_eq!(bar.await, Some(42));
95    /// #
96    /// # })
97    /// ```
98    pub struct Timeout<Fut: Future> {
99        #[pin]
100        future: Fut,
101        #[pin]
102        timer: Timer,
103    }
104}
105
106/* ┌────────────────────────────────────────────────────────────────────────────────────────────┐ *\
107 * │                                  trait TimeoutExt: Future                                  │ *
108\* └────────────────────────────────────────────────────────────────────────────────────────────┘ */
109
110/// An extension trait for [`Future`]s that provides a way to create [`Timeout`]s.
111pub trait TimeoutExt: Future {
112    /// Given a [`Duration`], creates and returns a new [`Timeout`] that will poll both the future
113    /// and a [`Timer`] that will complete after the provided duration, and return the future's
114    /// output or [`None`] if the timer completes first.
115    ///
116    /// ## Example
117    ///
118    /// ```rust
119    /// use async_io::Timer;
120    /// # use futures_lite::future;
121    /// use smol_timeout::TimeoutExt;
122    /// use std::time::Duration;
123    ///
124    /// # future::block_on(async {
125    /// #
126    /// let foo = async {
127    ///     Timer::after(Duration::from_millis(250)).await;
128    ///     24
129    /// };
130    ///
131    /// let foo = foo.timeout(Duration::from_millis(100));
132    /// assert_eq!(foo.await, None);
133    ///
134    /// let bar = async {
135    ///     Timer::after(Duration::from_millis(100)).await;
136    ///     42
137    /// };
138    ///
139    /// let bar = bar.timeout(Duration::from_millis(250));
140    /// assert_eq!(bar.await, Some(42));
141    /// #
142    /// # })
143    /// ```
144    fn timeout(self, after: Duration) -> Timeout<Self>
145    where
146        Self: Sized,
147    {
148        Timeout {
149            future: self,
150            timer: Timer::after(after),
151        }
152    }
153}
154
155impl<Fut: Future> TimeoutExt for Fut {}
156
157/* ┌────────────────────────────────────────────────────────────────────────────────────────────┐ *\
158 * │                                impl Future for Timeout<Fut>                                │ *
159\* └────────────────────────────────────────────────────────────────────────────────────────────┘ */
160
161impl<Fut: Future> Future for Timeout<Fut> {
162    type Output = Option<Fut::Output>;
163
164    fn poll(self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Self::Output> {
165        let this = self.project();
166
167        if this.timer.poll(ctx).is_ready() {
168            return Poll::Ready(None);
169        }
170
171        if let Poll::Ready(output) = this.future.poll(ctx) {
172            return Poll::Ready(Some(output));
173        }
174
175        Poll::Pending
176    }
177}