task: add AbortHandle type for cancelling tasks in a JoinSet

tokio/src/runtime/task/abort.rs

@@ -0,0 +1,69 @@
+use crate::runtime::task::RawTask;
+use std::fmt;
+use std::panic::{RefUnwindSafe, UnwindSafe};
+
+/// An owned permission to abort a spawned task, without awaiting its completion.
+///
+/// Unlike a [`JoinHandle`], an `AbortHandle` does *not* represent the
+/// permission to await the task's completion, only to terminate it.
+///
+/// The task may be aborted by calling the [`AbortHandle::abort`] method.
+/// Dropping an `AbortHandle` releases the permission to terminate the task
+/// --- it does *not* abort the task.
+///
+/// **Note**: This is an [unstable API][unstable]. The public API of this type
+/// may break in 1.x releases. See [the documentation on unstable
+/// features][unstable] for details.
+///
+/// [unstable]: crate#unstable-features
+/// [`JoinHandle`]: crate::task::JoinHandle
+#[cfg_attr(docsrs, doc(cfg(all(feature = "rt", tokio_unstable))))]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+pub struct AbortHandle {
+    raw: Option<RawTask>,
+}
+
+impl AbortHandle {
+    pub(super) fn new(raw: Option<RawTask>) -> Self {
+        Self { raw }
+    }
+
+    /// Abort the task associated with the handle.
+    ///
+    /// Awaiting a cancelled task might complete as usual if the task was
+    /// already completed at the time it was cancelled, but most likely it
+    /// will fail with a [cancelled] `JoinError`.
+    ///
+    /// If the task was already cancelled, such as by [`JoinHandle::abort`],
+    /// this method will do nothing.
+    ///
+    /// [cancelled]: method@super::error::JoinError::is_cancelled
+    // the `AbortHandle` type is only publicly exposed when `tokio_unstable` is
+    // enabled, but it is still defined for testing purposes.
+    #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+    pub fn abort(self) {
+        if let Some(raw) = self.raw {
+            raw.remote_abort();
+        }
+    }
+}
+
+unsafe impl Send for AbortHandle {}
+unsafe impl Sync for AbortHandle {}
+
+impl UnwindSafe for AbortHandle {}
+impl RefUnwindSafe for AbortHandle {}
+
+impl fmt::Debug for AbortHandle {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("AbortHandle").finish()
+    }
+}
+
+impl Drop for AbortHandle {
+    fn drop(&mut self) {
+        if let Some(raw) = self.raw.take() {
+            raw.drop_abort_handle();
+        }
+    }
+}

tokio/src/runtime/task/join.rs

@@ -210,6 +210,16 @@ impl<T> JoinHandle<T> {
             }
         }
     }
+
+    /// Returns a new `AbortHandle` that can be used to remotely abort this task.
+    #[cfg(any(tokio_unstable, test))]
+    pub(crate) fn abort_handle(&self) -> super::AbortHandle {
+        let raw = self.raw.map(|raw| {
+            raw.ref_inc();
+            raw
+        });
+        super::AbortHandle::new(raw)
+    }
 }
 
 impl<T> Unpin for JoinHandle<T> {}

tokio/src/runtime/task/mod.rs

@@ -155,7 +155,14 @@ cfg_rt_multi_thread! {
     pub(super) use self::inject::Inject;
 }
 
+#[cfg(all(feature = "rt", any(tokio_unstable, test)))]
+mod abort;
 mod join;
+
+#[cfg(all(feature = "rt", any(tokio_unstable, test)))]
+#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
+pub use self::abort::AbortHandle;
+
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
 pub use self::join::JoinHandle;
 

tokio/src/runtime/task/raw.rs

@@ -27,6 +27,9 @@ pub(super) struct Vtable {
     /// The join handle has been dropped.
     pub(super) drop_join_handle_slow: unsafe fn(NonNull<Header>),
 
+    /// An abort handle has been dropped.
+    pub(super) drop_abort_handle: unsafe fn(NonNull<Header>),
+
     /// The task is remotely aborted.
     pub(super) remote_abort: unsafe fn(NonNull<Header>),
 
@@ -42,6 +45,7 @@ pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable {
         try_read_output: try_read_output::<T, S>,
         try_set_join_waker: try_set_join_waker::<T, S>,
         drop_join_handle_slow: drop_join_handle_slow::<T, S>,
+        drop_abort_handle: drop_abort_handle::<T, S>,
         remote_abort: remote_abort::<T, S>,
         shutdown: shutdown::<T, S>,
     }
@@ -104,6 +108,11 @@ impl RawTask {
         unsafe { (vtable.drop_join_handle_slow)(self.ptr) }
     }
 
+    pub(super) fn drop_abort_handle(self) {
+        let vtable = self.header().vtable;
+        unsafe { (vtable.drop_abort_handle)(self.ptr) }
+    }
+
     pub(super) fn shutdown(self) {
         let vtable = self.header().vtable;
         unsafe { (vtable.shutdown)(self.ptr) }
@@ -113,6 +122,13 @@ impl RawTask {
         let vtable = self.header().vtable;
         unsafe { (vtable.remote_abort)(self.ptr) }
     }
+
+    /// Increment the task's reference count.
+    ///
+    /// Currently, this is used only when creating an `AbortHandle`.
+    pub(super) fn ref_inc(self) {
+        self.header().state.ref_inc();
+    }
 }
 
 impl Clone for RawTask {
@@ -154,6 +170,11 @@ unsafe fn drop_join_handle_slow<T: Future, S: Schedule>(ptr: NonNull<Header>) {
     harness.drop_join_handle_slow()
 }
 
+unsafe fn drop_abort_handle<T: Future, S: Schedule>(ptr: NonNull<Header>) {
+    let harness = Harness::<T, S>::from_raw(ptr);
+    harness.drop_reference();
+}
+
 unsafe fn remote_abort<T: Future, S: Schedule>(ptr: NonNull<Header>) {
     let harness = Harness::<T, S>::from_raw(ptr);
     harness.remote_abort()

tokio/src/runtime/tests/task.rs

@@ -78,6 +78,44 @@ fn create_drop2() {
     handle.assert_dropped();
 }
 
+#[test]
+fn drop_abort_handle1() {
+    let (ad, handle) = AssertDrop::new();
+    let (notified, join) = unowned(
+        async {
+            drop(ad);
+            unreachable!()
+        },
+        NoopSchedule,
+    );
+    let abort = join.abort_handle();
+    drop(join);
+    handle.assert_not_dropped();
+    drop(notified);
+    handle.assert_not_dropped();
+    drop(abort);
+    handle.assert_dropped();
+}
+
+#[test]
+fn drop_abort_handle2() {
+    let (ad, handle) = AssertDrop::new();
+    let (notified, join) = unowned(
+        async {
+            drop(ad);
+            unreachable!()
+        },
+        NoopSchedule,
+    );
+    let abort = join.abort_handle();
+    drop(notified);
+    handle.assert_not_dropped();
+    drop(abort);
+    handle.assert_not_dropped();
+    drop(join);
+    handle.assert_dropped();
+}
+
 // Shutting down through Notified works
 #[test]
 fn create_shutdown1() {