Rust abortable Future

abortable.rs

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;

#[tokio::main]
async fn main() {
    struct Abortable<T, Fut: Future<Output = T> + Unpin> {
        aborted: bool,
        future: Fut,
    }

    impl<T, Fut: Future<Output = T> + Unpin> Abortable<T, Fut> {
        pub fn new(future: Fut) -> Self {
            Self {
                aborted: false,
                future,
            }
        }

        pub fn abort(&mut self) {
            self.aborted = true;
        }
    }

    impl<T, Fut: Future<Output = T> + Unpin> Future for Abortable<T, Fut> {
        type Output = Option<T>;

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            if self.aborted {
                return Poll::Ready(None);
            }

            match Pin::new(&mut self.future).poll(cx) {
                Poll::Ready(r) => Poll::Ready(Some(r)),
                Poll::Pending => Poll::Pending,
            }
        }
    }

    let fut1 = Abortable::new(Box::pin(tokio::time::sleep(Duration::from_millis(10))));
    assert_eq!(fut1.await, Some(()));
    
    let mut fut2 = Abortable::new(Box::pin(tokio::time::sleep(Duration::from_millis(10))));
    fut2.abort();
    assert_eq!(fut2.await, None);
}