azure_identity/
timeout.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
// Copyright (c) 2020 Yoshua Wuyts
//
// based on https://crates.io/crates/futures-time
// Licensed under either of Apache License, Version 2.0 or MIT license at your option.

use azure_core::sleep::{sleep, Sleep};
use futures::Future;
use std::time::Duration;
use std::{
    pin::Pin,
    task::{Context, Poll},
};

#[pin_project::pin_project]
#[derive(Debug)]
pub(crate) struct Timeout<F, D> {
    #[pin]
    future: F,
    #[pin]
    deadline: D,
    completed: bool,
}

impl<F, D> Timeout<F, D> {
    pub(crate) fn new(future: F, deadline: D) -> Self {
        Self {
            future,
            deadline,
            completed: false,
        }
    }
}

impl<F: Future, D: Future> Future for Timeout<F, D> {
    type Output = azure_core::Result<F::Output>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();

        assert!(!*this.completed, "future polled after completing");

        match this.future.poll(cx) {
            Poll::Ready(v) => {
                *this.completed = true;
                Poll::Ready(Ok(v))
            }
            Poll::Pending => match this.deadline.poll(cx) {
                Poll::Ready(_) => {
                    *this.completed = true;
                    Poll::Ready(Err(azure_core::error::Error::with_message(
                        azure_core::error::ErrorKind::Other,
                        || String::from("operation timed out"),
                    )))
                }
                Poll::Pending => Poll::Pending,
            },
        }
    }
}

pub(crate) trait TimeoutExt: Future {
    fn timeout(self, duration: Duration) -> Timeout<Self, Sleep>
    where
        Self: Sized,
    {
        Timeout::new(self, sleep(duration))
    }
}

impl<T> TimeoutExt for T where T: Future {}