azure_identity/
timeout.rs

1// Copyright (c) 2020 Yoshua Wuyts
2//
3// based on https://crates.io/crates/futures-time
4// Licensed under either of Apache License, Version 2.0 or MIT license at your option.
5
6use azure_core::sleep::{sleep, Sleep};
7use futures::Future;
8use std::time::Duration;
9use std::{
10    pin::Pin,
11    task::{Context, Poll},
12};
13
14#[pin_project::pin_project]
15#[derive(Debug)]
16pub(crate) struct Timeout<F, D> {
17    #[pin]
18    future: F,
19    #[pin]
20    deadline: D,
21    completed: bool,
22}
23
24impl<F, D> Timeout<F, D> {
25    pub(crate) fn new(future: F, deadline: D) -> Self {
26        Self {
27            future,
28            deadline,
29            completed: false,
30        }
31    }
32}
33
34impl<F: Future, D: Future> Future for Timeout<F, D> {
35    type Output = azure_core::Result<F::Output>;
36
37    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
38        let this = self.project();
39
40        assert!(!*this.completed, "future polled after completing");
41
42        match this.future.poll(cx) {
43            Poll::Ready(v) => {
44                *this.completed = true;
45                Poll::Ready(Ok(v))
46            }
47            Poll::Pending => match this.deadline.poll(cx) {
48                Poll::Ready(_) => {
49                    *this.completed = true;
50                    Poll::Ready(Err(azure_core::error::Error::with_message(
51                        azure_core::error::ErrorKind::Other,
52                        || String::from("operation timed out"),
53                    )))
54                }
55                Poll::Pending => Poll::Pending,
56            },
57        }
58    }
59}
60
61pub(crate) trait TimeoutExt: Future {
62    fn timeout(self, duration: Duration) -> Timeout<Self, Sleep>
63    where
64        Self: Sized,
65    {
66        Timeout::new(self, sleep(duration))
67    }
68}
69
70impl<T> TimeoutExt for T where T: Future {}