async_shutdown/
wrap_delay_shutdown.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4
5use crate::DelayShutdownToken;
6
7/// Wrapped future that delays shutdown completion until it completes or until it is droppped.
8#[must_use = "futures must be polled to make progress"]
9pub struct WrapDelayShutdown<T: Clone, F> {
10	pub(crate) delay_token: Option<DelayShutdownToken<T>>,
11	pub(crate) future: F,
12}
13
14impl<T: Clone, F: Future> Future for WrapDelayShutdown<T, F> {
15	type Output = F::Output;
16
17	#[inline]
18	fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
19		// SAFETY: We never move `future`, so we can not violate the requirements of `F`.
20		unsafe {
21			let me = self.get_unchecked_mut();
22			match Pin::new_unchecked(&mut me.future).poll(context) {
23				Poll::Pending => Poll::Pending,
24				Poll::Ready(value) => {
25					me.delay_token = None;
26					Poll::Ready(value)
27				},
28			}
29		}
30	}
31}