Bug 1746533 - Part 2: Use the fallible Poll implementation in moz_task, r=emilio

This should avoid crash issues caused by background task queues being destroyed
and dropping their Runnables.

Differential Revision: https://phabricator.services.mozilla.com/D135410

xpcom/rust/moz_task/src/executor.rs

@@ -3,13 +3,73 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
 use crate::{get_current_thread, DispatchOptions, RunnableBuilder};
-use async_task::Task;
-use std::cell::Cell;
-use std::sync::Arc;
-use std::{fmt::Debug, future::Future, ptr};
+use std::{
+    cell::Cell,
+    fmt::Debug,
+    future::Future,
+    pin::Pin,
+    ptr,
+    sync::Arc,
+    task::{Context, Poll},
+};
 use xpcom::interfaces::{nsIEventTarget, nsIRunnablePriority};
 use xpcom::RefPtr;
 
+/// A spawned task.
+///
+/// A [`AsyncTask`] can be awaited to retrieve the output of its future.
+///
+/// Dropping an [`AsyncTask`] cancels it, which means its future won't be polled
+/// again. To drop the [`AsyncTask`] handle without canceling it, use
+/// [`detach()`][`AsyncTask::detach()`] instead. To cancel a task gracefully and
+/// wait until it is fully destroyed, use the [`cancel()`][AsyncTask::cancel()]
+/// method.
+///
+/// A task which is cancelled due to the nsIEventTarget it was dispatched to no
+/// longer accepting events will never be resolved.
+#[derive(Debug)]
+#[must_use = "tasks get canceled when dropped, use `.detach()` to run them in the background"]
+pub struct AsyncTask<T> {
+    task: async_task::FallibleTask<T>,
+}
+
+impl<T> AsyncTask<T> {
+    fn new(task: async_task::Task<T>) -> Self {
+        AsyncTask {
+            task: task.fallible(),
+        }
+    }
+
+    /// Detaches the task to let it keep running in the background.
+    pub fn detach(self) {
+        self.task.detach()
+    }
+
+    /// Cancels the task and waits for it to stop running.
+    ///
+    /// Returns the task's output if it was completed just before it got canceled, or [`None`] if
+    /// it didn't complete.
+    ///
+    /// While it's possible to simply drop the [`Task`] to cancel it, this is a cleaner way of
+    /// canceling because it also waits for the task to stop running.
+    pub async fn cancel(self) -> Option<T> {
+        self.task.cancel().await
+    }
+}
+
+impl<T> Future for AsyncTask<T> {
+    type Output = T;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        // Wrap the future produced by `AsyncTask` to never resolve if the
+        // Runnable was dropped, and the task was cancelled.
+        match Pin::new(&mut self.task).poll(cx) {
+            Poll::Ready(Some(t)) => Poll::Ready(t),
+            Poll::Ready(None) | Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
 enum SpawnTarget {
     BackgroundTask,
     EventTarget(RefPtr<nsIEventTarget>),
@@ -119,7 +179,7 @@ where
     F::Output: Send + 'static,
 {
     /// Run the future on the background task pool.
-    pub fn spawn(self) -> Task<F::Output> {
+    pub fn spawn(self) -> AsyncTask<F::Output> {
         let config = Arc::new(TaskSpawnConfig {
             name: self.name,
             priority: self.priority,
@@ -130,11 +190,11 @@ where
             schedule(config.clone(), runnable)
         });
         runnable.schedule();
-        task
+        AsyncTask::new(task)
     }
 
     /// Run the future on the specified nsIEventTarget.
-    pub fn spawn_onto(self, target: &nsIEventTarget) -> Task<F::Output> {
+    pub fn spawn_onto(self, target: &nsIEventTarget) -> AsyncTask<F::Output> {
         let config = Arc::new(TaskSpawnConfig {
             name: self.name,
             priority: self.priority,
@@ -145,7 +205,7 @@ where
             schedule(config.clone(), runnable)
         });
         runnable.schedule();
-        task
+        AsyncTask::new(task)
     }
 }
 
@@ -163,7 +223,7 @@ where
     /// This method may panic if run on a thread which cannot run local futures
     /// (e.g. due to it is not being an XPCOM thread, or if we are very late
     /// during shutdown).
-    pub fn spawn_local(self) -> Task<F::Output> {
+    pub fn spawn_local(self) -> AsyncTask<F::Output> {
         let current_thread = get_current_thread().expect("cannot get current thread");
         let config = Arc::new(TaskSpawnConfig {
             name: self.name,
@@ -175,13 +235,13 @@ where
             schedule(config.clone(), runnable)
         });
         runnable.schedule();
-        task
+        AsyncTask::new(task)
     }
 }
 
 /// Spawn a future onto the background task pool. The future will not be run on
 /// the main thread.
-pub fn spawn<F>(name: &'static str, future: F) -> Task<F::Output>
+pub fn spawn<F>(name: &'static str, future: F) -> AsyncTask<F::Output>
 where
     F: Future + Send + 'static,
     F::Output: Send + 'static,
@@ -191,7 +251,7 @@ where
 
 /// Spawn a potentially-blocking future onto the background task pool. The
 /// future will not be run on the main thread.
-pub fn spawn_blocking<F>(name: &'static str, future: F) -> Task<F::Output>
+pub fn spawn_blocking<F>(name: &'static str, future: F) -> AsyncTask<F::Output>
 where
     F: Future + Send + 'static,
     F::Output: Send + 'static,
@@ -200,7 +260,7 @@ where
 }
 
 /// Spawn a local future onto the current thread.
-pub fn spawn_local<F>(name: &'static str, future: F) -> Task<F::Output>
+pub fn spawn_local<F>(name: &'static str, future: F) -> AsyncTask<F::Output>
 where
     F: Future + 'static,
     F::Output: 'static,
@@ -208,7 +268,7 @@ where
     TaskBuilder::new(name, future).spawn_local()
 }
 
-pub fn spawn_onto<F>(name: &'static str, target: &nsIEventTarget, future: F) -> Task<F::Output>
+pub fn spawn_onto<F>(name: &'static str, target: &nsIEventTarget, future: F) -> AsyncTask<F::Output>
 where
     F: Future + Send + 'static,
     F::Output: Send + 'static,
@@ -220,7 +280,7 @@ pub fn spawn_onto_blocking<F>(
     name: &'static str,
     target: &nsIEventTarget,
     future: F,
-) -> Task<F::Output>
+) -> AsyncTask<F::Output>
 where
     F: Future + Send + 'static,
     F::Output: Send + 'static,

xpcom/rust/moz_task/src/lib.rs

@@ -12,13 +12,9 @@ pub use dispatcher::{dispatch_background_task, dispatch_local, dispatch_onto, Ru
 mod event_loop;
 mod executor;
 pub use executor::{
-    spawn, spawn_blocking, spawn_local, spawn_onto, spawn_onto_blocking, TaskBuilder,
+    spawn, spawn_blocking, spawn_local, spawn_onto, spawn_onto_blocking, AsyncTask, TaskBuilder,
 };
 
-// FIXME: Unfortunately directly re-exporting as `Task` conflicts with the task
-// trait below. This type is useful for folks using the `spawn*` methods.
-pub use async_task::Task as AsyncTask;
-
 // Expose functions intended to be used only in gtest via this module.
 // We don't use a feature gate here to stop the need to compile all crates that
 // depend upon `moz_task` twice.