async_shutdown/
wrap_trigger_shutdown.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4
5use crate::TriggerShutdownToken;
6
7/// Wrapped future that triggers a shutdown when it completes or when it is dropped.
8#[must_use = "futures must be polled to make progress"]
9pub struct WrapTriggerShutdown<T: Clone, F> {
10	pub(crate) trigger_shutdown_token: Option<TriggerShutdownToken<T>>,
11	pub(crate) future: F,
12}
13
14impl<T: Clone, F: Future> Future for WrapTriggerShutdown<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.trigger_shutdown_token = None;
26					Poll::Ready(value)
27				},
28			}
29		}
30	}
31}