steamkit/util/
task_extensions.rs

1//! Ignoring cancellation errors once a token is cancelled.
2
3use std::fmt;
4use std::future::Future;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8use tokio_util::sync::CancellationToken;
9
10use crate::error::SteamError;
11
12/// A future that swallows its cancellation error if `token` was cancelled.
13pub struct IgnoringCancellation<F> {
14    /// The wrapped future.
15    inner: Pin<Box<F>>,
16    /// Decides whether a cancellation error is swallowed.
17    token: CancellationToken,
18}
19
20impl<F> IgnoringCancellation<F> {
21    /// Wraps `inner` so its cancellation error is ignored once `token` is
22    /// cancelled.
23    pub fn new(inner: F, token: CancellationToken) -> IgnoringCancellation<F> {
24        IgnoringCancellation {
25            inner: Box::pin(inner),
26            token,
27        }
28    }
29}
30
31impl<F> fmt::Debug for IgnoringCancellation<F> {
32    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
33        formatter
34            .debug_struct("IgnoringCancellation")
35            .field("cancelled", &self.token.is_cancelled())
36            .finish()
37    }
38}
39
40impl<F> Future for IgnoringCancellation<F>
41where
42    F: Future<Output = Result<(), SteamError>>,
43{
44    type Output = Result<(), SteamError>;
45
46    fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
47        let this = self.get_mut();
48
49        match this.inner.as_mut().poll(context) {
50            Poll::Pending => Poll::Pending,
51            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
52            Poll::Ready(Err(error)) => {
53                if matches!(error, SteamError::Cancelled) && this.token.is_cancelled() {
54                    Poll::Ready(Ok(()))
55                } else {
56                    Poll::Ready(Err(error))
57                }
58            }
59        }
60    }
61}
62
63/// Awaits `task`, ignoring its cancellation error once `token` is cancelled.
64/// The trait version is [`TaskExtensions::ignoring_cancellation`].
65pub fn ignoring_cancellation<F>(task: F, token: CancellationToken) -> IgnoringCancellation<F>
66where
67    F: Future<Output = Result<(), SteamError>>,
68{
69    IgnoringCancellation::new(task, token)
70}
71
72/// Extension methods for futures.
73pub trait TaskExtensions: Future<Output = Result<(), SteamError>> + Sized {
74    /// Awaits this future, ignoring its cancellation error once `token` is
75    /// cancelled. There's also the free function [`ignoring_cancellation`].
76    fn ignoring_cancellation(self, token: CancellationToken) -> IgnoringCancellation<Self> {
77        IgnoringCancellation::new(self, token)
78    }
79}
80
81impl<F> TaskExtensions for F where F: Future<Output = Result<(), SteamError>> + Sized {}
steamkit/util/task_extensions.rs

steamkit/util/
task_extensions.rs
