async_shutdown/
wrap_cancel.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4
5use crate::shutdown_signal::ShutdownSignal;
6
7/// Wrapped future that is automatically cancelled when a shutdown is triggered.
8///
9/// If the wrapped future completes before the shutdown is triggered,
10/// the output of the original future is yielded as `Ok(value)`.
11///
12/// If the shutdown is triggered before the wrapped future completes,
13/// the original future is dropped and the shutdown reason is yielded as `Err(shutdown_reason)`.
14#[must_use = "futures must be polled to make progress"]
15pub struct WrapCancel<T: Clone, F> {
16	pub(crate) shutdown_signal: ShutdownSignal<T>,
17	pub(crate) future: Result<F, T>,
18}
19
20impl<T: Clone, F: Future> Future for WrapCancel<T, F> {
21	type Output = Result<F::Output, T>;
22
23	#[inline]
24	fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
25		// SAFETY: We never move `future`, so we can not violate the requirements of `F`.
26		// We do drop it, but that's allowed by `Pin`.
27		let me = unsafe { self.get_unchecked_mut() };
28
29		match &mut me.future {
30			Err(e) => return Poll::Ready(Err(e.clone())),
31			Ok(future) => {
32				let future = unsafe { Pin::new_unchecked(future) };
33				if let Poll::Ready(value) = future.poll(context) {
34					return Poll::Ready(Ok(value));
35				}
36			},
37		}
38
39		// Otherwise check if the shutdown signal has been given.
40		let shutdown = Pin::new(&mut me.shutdown_signal).poll(context);
41		match shutdown {
42			Poll::Ready(reason) => {
43				me.future = Err(reason.clone());
44				Poll::Ready(Err(reason))
45			},
46			Poll::Pending => Poll::Pending,
47		}
48	}
49}
async_shutdown/wrap_cancel.rs

async_shutdown/
wrap_cancel.rs
