aws_smithy_async/rt/
sleep.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Provides an [`AsyncSleep`] trait that returns a future that sleeps for a given duration,
7//! and implementations of `AsyncSleep` for different async runtimes.
8
9use std::fmt::{Debug, Formatter};
10use std::future::Future;
11use std::pin::Pin;
12use std::sync::Arc;
13use std::task::{Context, Poll};
14use std::time::Duration;
15
16/// Async trait with a `sleep` function.
17pub trait AsyncSleep: Debug + Send + Sync {
18    /// Returns a future that sleeps for the given `duration` of time.
19    fn sleep(&self, duration: Duration) -> Sleep;
20}
21
22impl<T> AsyncSleep for Box<T>
23where
24    T: AsyncSleep,
25    T: ?Sized,
26{
27    fn sleep(&self, duration: Duration) -> Sleep {
28        T::sleep(self, duration)
29    }
30}
31
32impl<T> AsyncSleep for Arc<T>
33where
34    T: AsyncSleep,
35    T: ?Sized,
36{
37    fn sleep(&self, duration: Duration) -> Sleep {
38        T::sleep(self, duration)
39    }
40}
41
42/// Wrapper type for sharable `AsyncSleep`
43#[derive(Clone, Debug)]
44pub struct SharedAsyncSleep(Arc<dyn AsyncSleep>);
45
46impl SharedAsyncSleep {
47    /// Create a new `SharedAsyncSleep` from `AsyncSleep`
48    pub fn new(sleep: impl AsyncSleep + 'static) -> Self {
49        Self(Arc::new(sleep))
50    }
51}
52
53impl AsRef<dyn AsyncSleep> for SharedAsyncSleep {
54    fn as_ref(&self) -> &(dyn AsyncSleep + 'static) {
55        self.0.as_ref()
56    }
57}
58
59impl From<Arc<dyn AsyncSleep>> for SharedAsyncSleep {
60    fn from(sleep: Arc<dyn AsyncSleep>) -> Self {
61        SharedAsyncSleep(sleep)
62    }
63}
64
65impl AsyncSleep for SharedAsyncSleep {
66    fn sleep(&self, duration: Duration) -> Sleep {
67        self.0.sleep(duration)
68    }
69}
70
71#[cfg(feature = "rt-tokio")]
72/// Returns a default sleep implementation based on the features enabled
73pub fn default_async_sleep() -> Option<SharedAsyncSleep> {
74    Some(SharedAsyncSleep::from(sleep_tokio()))
75}
76
77#[cfg(not(feature = "rt-tokio"))]
78/// Returns a default sleep implementation based on the features enabled
79pub fn default_async_sleep() -> Option<SharedAsyncSleep> {
80    None
81}
82
83/// Future returned by [`AsyncSleep`].
84#[non_exhaustive]
85#[must_use]
86pub struct Sleep(Pin<Box<dyn Future<Output = ()> + Send + Sync + 'static>>);
87
88impl Debug for Sleep {
89    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
90        write!(f, "Sleep")
91    }
92}
93
94impl Sleep {
95    /// Create a new [`Sleep`] future
96    ///
97    /// The provided future will be Boxed.
98    pub fn new(future: impl Future<Output = ()> + Send + Sync + 'static) -> Sleep {
99        Sleep(Box::pin(future))
100    }
101}
102
103impl Future for Sleep {
104    type Output = ();
105
106    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
107        self.0.as_mut().poll(cx)
108    }
109}
110
111/// Implementation of [`AsyncSleep`] for Tokio.
112#[non_exhaustive]
113#[cfg(feature = "rt-tokio")]
114#[derive(Debug, Default)]
115pub struct TokioSleep;
116
117#[cfg(feature = "rt-tokio")]
118impl TokioSleep {
119    /// Create a new [`AsyncSleep`] implementation using the Tokio hashed wheel sleep implementation
120    pub fn new() -> TokioSleep {
121        Default::default()
122    }
123}
124
125#[cfg(feature = "rt-tokio")]
126impl AsyncSleep for TokioSleep {
127    fn sleep(&self, duration: Duration) -> Sleep {
128        Sleep::new(tokio::time::sleep(duration))
129    }
130}
131
132#[cfg(feature = "rt-tokio")]
133fn sleep_tokio() -> Arc<dyn AsyncSleep> {
134    Arc::new(TokioSleep::new())
135}