diff --git a/runtime/src/task.rs b/runtime/src/task.rs index b75aca89eb..72f408e036 100644 --- a/runtime/src/task.rs +++ b/runtime/src/task.rs @@ -168,9 +168,21 @@ impl Task { Some(stream) => { let (stream, handle) = stream::abortable(stream); - (Self(Some(boxed_stream(stream))), Handle(Some(handle))) + ( + Self(Some(boxed_stream(stream))), + Handle { + raw: Some(handle), + abort_on_drop: false, + }, + ) } - None => (Self(None), Handle(None)), + None => ( + Self(None), + Handle { + raw: None, + abort_on_drop: false, + }, + ), } } @@ -195,19 +207,34 @@ impl Task { /// A handle to a [`Task`] that can be used for aborting it. #[derive(Debug, Clone)] -pub struct Handle(Option); +pub struct Handle { + raw: Option, + abort_on_drop: bool, +} impl Handle { /// Aborts the [`Task`] of this [`Handle`]. pub fn abort(&self) { - if let Some(handle) = &self.0 { + if let Some(handle) = &self.raw { handle.abort(); } } + /// Returns a new [`Handle`] that will call [`Handle::abort`] whenever + /// it is dropped. + /// + /// This can be really useful if you do not want to worry about calling + /// [`Handle::abort`] yourself. + pub fn abort_on_drop(mut self) -> Self { + Self { + raw: self.raw.take(), + abort_on_drop: true, + } + } + /// Returns `true` if the [`Task`] of this [`Handle`] has been aborted. pub fn is_aborted(&self) -> bool { - if let Some(handle) = &self.0 { + if let Some(handle) = &self.raw { handle.is_aborted() } else { true @@ -215,6 +242,14 @@ impl Handle { } } +impl Drop for Handle { + fn drop(&mut self) { + if self.abort_on_drop { + self.abort(); + } + } +} + impl Task> { /// Executes a new [`Task`] after this one, only when it produces `Some` value. ///