Make SpawnHandle shutdown on drop by default (#36297)

This prevents a common class of bugs on startup where a spawn handle can be created,
but then the initialization process hits an error and the task gets leaked.

GitOrigin-RevId: 31a4b3f4dc37bce00bee1379d3e9e5bdfb866a25

Author: goffrie

crates/application/src/cron_jobs/mod.rs

@@ -231,7 +231,8 @@ impl<RT: Runtime> CronJobExecutor<RT> {
             );
             let context = self.clone();
             let tx = job_finished_tx.clone();
-            self.rt.spawn(
+            // TODO: cancel this handle with the application
+            self.rt.spawn_background(
                 "spawn_cron_job",
                 async move {
                     select_biased! {

crates/application/src/lib.rs

@@ -1220,7 +1220,8 @@ impl<RT: Runtime> Application<RT> {
             // Spawn running the action in a separate future. This way, even if we
             // get cancelled, it will continue to run to completion.
             let (tx, rx) = oneshot::channel();
-            self.runtime.spawn("run_action", async move {
+            // TODO: cancel this handle with the application
+            self.runtime.spawn_background("run_action", async move {
                 let result = run_action.await;
                 // Don't log errors if the caller has gone away.
                 _ = tx.send(result);
@@ -1278,24 +1279,26 @@ impl<RT: Runtime> Application<RT> {
             .map(|ctx| Span::root(format!("{}::http_actions_future", func_path!()), ctx))
             .unwrap_or(Span::noop());
         let response_streamer_ = response_streamer.clone();
-        self.runtime.spawn("run_http_action", async move {
-            let result = runner
-                .run_http_action(
-                    request_id,
-                    http_request,
-                    response_streamer_,
-                    identity,
-                    caller,
-                )
-                .in_span(span)
-                .await;
-            if let Err(Err(mut e)) = tx.send(result) {
-                // If the caller has gone away, and the result is a system error,
-                // log to sentry.
-                report_error(&mut e).await;
-            }
-            rt_.pause_client().wait("end_run_http_action").await;
-        });
+        // TODO: cancel this handle with the application
+        self.runtime
+            .spawn_background("run_http_action", async move {
+                let result = runner
+                    .run_http_action(
+                        request_id,
+                        http_request,
+                        response_streamer_,
+                        identity,
+                        caller,
+                    )
+                    .in_span(span)
+                    .await;
+                if let Err(Err(mut e)) = tx.send(result) {
+                    // If the caller has gone away, and the result is a system error,
+                    // log to sentry.
+                    report_error(&mut e).await;
+                }
+                rt_.pause_client().wait("end_run_http_action").await;
+            });
         let result = rx
             .await
             .context("run_http_action one shot sender dropped prematurely?")?;

crates/application/src/scheduled_jobs/mod.rs

@@ -328,7 +328,8 @@ impl<RT: Runtime> ScheduledJobExecutor<RT> {
                 &mut self.rt.rng(),
                 BTreeMap::new(),
             );
-            self.rt.spawn(
+            // TODO: cancel this handle with the application
+            self.rt.spawn_background(
                 "spawn_scheduled_job",
                 async move {
                     context.execute_job(job, job_id).await;

crates/aws_s3/src/storage.rs

@@ -672,19 +672,21 @@ impl<RT: Runtime> Drop for S3Upload<RT> {
     fn drop(&mut self) {
         if self.needs_abort_on_drop {
             let fut = self._abort();
-            self.runtime.spawn("abort_multipart_upload", async move {
-                if let Err(e) = fut.await {
-                    // abort-multipart-upload is idempotent. It has the following properties.
-                    //
-                    // abort after a successful abort - succeeds
-                    // abort after a successful complete - succeeds
-                    // complete after a successful abort - fails with a descriptive error.
-                    report_error(
-                        &mut anyhow::anyhow!(e).context("Couldn't async abort multipart upload"),
-                    )
-                    .await;
-                }
-            });
+            self.runtime
+                .spawn_background("abort_multipart_upload", async move {
+                    if let Err(e) = fut.await {
+                        // abort-multipart-upload is idempotent. It has the following properties.
+                        //
+                        // abort after a successful abort - succeeds
+                        // abort after a successful complete - succeeds
+                        // complete after a successful abort - fails with a descriptive error.
+                        report_error(
+                            &mut anyhow::anyhow!(e)
+                                .context("Couldn't async abort multipart upload"),
+                        )
+                        .await;
+                    }
+                });
         }
     }
 }

crates/common/src/runtime/mod.rs

@@ -96,9 +96,18 @@ impl From<tokio::task::JoinError> for JoinError {
     }
 }
 
+#[must_use = "Tasks are canceled when their `SpawnHandle` is dropped."]
 pub trait SpawnHandle: Send + Sync {
+    /// Stops the spawned task "soon". This happens asynchronously.
     fn shutdown(&mut self);
+    /// Wait for the spawned task to finish. Don't use this function directly,
+    /// call `.join()` instead.
     fn poll_join(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), JoinError>>;
+    /// Allows the spawned task to keep running indefinitely. By default, a task
+    /// is shut down on drop.
+    fn detach(self: Box<Self>);
+    /// Wait for the spawned task to finish. Returns an error if the task was
+    /// canceled (using `.shutdown()`) or panicked.
     fn join<'a>(mut self) -> impl Future<Output = Result<(), JoinError>> + 'a
     where
         Self: Sized + 'a,
@@ -115,6 +124,10 @@ impl<T: SpawnHandle + ?Sized> SpawnHandle for Box<T> {
     fn poll_join(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), JoinError>> {
         (**self).poll_join(cx)
     }
+
+    fn detach(self: Box<Self>) {
+        (*self).detach()
+    }
 }
 
 impl dyn SpawnHandle {
@@ -123,27 +136,61 @@ impl dyn SpawnHandle {
     pub fn join(self: Box<Self>) -> impl Future<Output = Result<(), JoinError>> {
         SpawnHandle::join(self)
     }
+
+    /// Wait for the spawn task to finish, but if the returned future is
+    /// canceled, the spawned task continues running as though it were
+    /// `detach()`ed.
+    pub fn join_or_detach(self: Box<Self>) -> impl Future<Output = Result<(), JoinError>> {
+        struct DetachOnDrop(Option<Box<dyn SpawnHandle>>);
+        impl Drop for DetachOnDrop {
+            fn drop(&mut self) {
+                self.0.take().expect("lost spawn handle?").detach();
+            }
+        }
+        let mut handle = DetachOnDrop(Some(self));
+        future::poll_fn(move |cx| handle.0.as_mut().expect("lost spawn handle?").poll_join(cx))
+    }
+
+    pub fn shutdown_and_join(self: Box<Self>) -> impl Future<Output = anyhow::Result<()>> {
+        shutdown_and_join(self)
+    }
 }
 
 pub struct TokioSpawnHandle {
-    handle: tokio::task::JoinHandle<()>,
+    handle: Option<tokio::task::JoinHandle<()>>,
 }
 
 impl From<tokio::task::JoinHandle<()>> for TokioSpawnHandle {
     fn from(handle: tokio::task::JoinHandle<()>) -> Self {
-        Self { handle }
+        Self {
+            handle: Some(handle),
+        }
     }
 }
 
 impl SpawnHandle for TokioSpawnHandle {
     fn shutdown(&mut self) {
-        self.handle.abort();
+        self.handle.as_ref().expect("shutdown after detach").abort();
     }
 
     fn poll_join(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), JoinError>> {
-        std::task::ready!(Pin::new(&mut self.handle).poll(cx))?;
+        std::task::ready!(
+            Pin::new(&mut self.handle.as_mut().expect("poll after detach")).poll(cx)
+        )?;
         Poll::Ready(Ok(()))
     }
+
+    fn detach(mut self: Box<Self>) {
+        self.handle.take();
+    }
+}
+
+impl Drop for TokioSpawnHandle {
+    fn drop(&mut self) {
+        if let Some(handle) = &self.handle {
+            handle.abort();
+        }
+    }
 }
 
 /// Shutdown the associated future, preempting it at its next yield point, and
@@ -230,16 +277,29 @@ pub trait Runtime: Clone + Sync + Send + 'static {
     fn wait(&self, duration: Duration) -> Pin<Box<dyn FusedFuture<Output = ()> + Send + 'static>>;
 
     /// Spawn a future on the runtime's executor.
+    ///
+    /// The spawned task will be canceled if the returned `SpawnHandle` is
+    /// dropped, unless `detach()` is called on it.
     fn spawn(
         &self,
         name: &'static str,
         f: impl Future<Output = ()> + Send + 'static,
     ) -> Box<dyn SpawnHandle>;
 
+    /// Shorthand for `spawn().detach()`
+    ///
+    /// This should only be used for tasks that are best-effort (e.g. cleaning
+    /// up partial progress) or that are truly process-global.
+    fn spawn_background(&self, name: &'static str, f: impl Future<Output = ()> + Send + 'static) {
+        self.spawn(name, f).detach()
+    }
+
     /// Spawn a future on a reserved OS thread. This is only really necessary
     /// for libraries like `V8` that care about being called from a
     /// particular thread.
-    #[must_use = "Threads are canceled when their `SpawnHandle` is dropped."]
+    ///
+    /// The spawned task will be canceled if the returned `SpawnHandle` is
+    /// dropped, unless `detach()` is called on it.
     fn spawn_thread<Fut: Future<Output = ()>, F: FnOnce() -> Fut + Send + 'static>(
         &self,
         name: &str,

crates/isolate/src/tests/fetch.rs

@@ -215,7 +215,7 @@ async fn test_fetch_basic(rt: ProdRuntime) -> anyhow::Result<()> {
                     .expect("invalid response")
             }),
         );
-    rt.spawn("test_server", serve(router, 4545));
+    let _server = rt.spawn("test_server", serve(router, 4545));
     let redirected_router = Router::new().route(
         "/print_auth",
         get(|req: Request<Body>| async move {
@@ -232,7 +232,7 @@ async fn test_fetch_basic(rt: ProdRuntime) -> anyhow::Result<()> {
                 .expect("invalid response")
         }),
     );
-    rt.spawn("test_router", serve(redirected_router, 4547));
+    let _router = rt.spawn("test_router", serve(redirected_router, 4547));
 
     let t = UdfTest::default(rt).await?;
     must_let!(let (ConvexValue::String(r), _outcome, log_lines) = t.action_outcome_and_log_lines(
@@ -301,7 +301,7 @@ async fn test_fetch_timing(rt: ProdRuntime) -> anyhow::Result<()> {
                     .expect("invalid response")
             }),
         );
-    rt.spawn("test_router", serve(router, 4546));
+    let _router = rt.spawn("test_router", serve(router, 4546));
 
     let t = UdfTest::default(rt.clone()).await?;
 
@@ -354,7 +354,7 @@ async fn test_fetch_abort(rt: ProdRuntime) -> anyhow::Result<()> {
             Response::builder().body(Body::from("ok")).unwrap()
         }),
     );
-    rt.spawn("test_router", serve(router, 4548));
+    let _router = rt.spawn("test_router", serve(router, 4548));
 
     // fetchAbort is an action that fetches from /pause, and in parallel it
     // waits for triggerAbort. When triggerAbort is called, it aborts the fetch

crates/load_generator/src/main.rs

@@ -451,9 +451,9 @@ fn main() -> Result<(), MainError> {
         .map(|addr| register_prometheus_exporter(runtime.clone(), addr));
     let load_generator = async move {
         run(&config).await?;
-        if let Some((mut handle, flush)) = maybe_metrics {
+        if let Some((handle, flush)) = maybe_metrics {
             flush().await;
-            handle.shutdown();
+            handle.shutdown_and_join().await?;
         }
         Ok::<_, MainError>(())
     };

crates/local_backend/src/lib.rs

@@ -243,7 +243,7 @@ pub async fn make_app(
             config.beacon_tag.clone(),
             config.beacon_fields.clone(),
         );
-        runtime.spawn("beacon_worker", beacon_future);
+        runtime.spawn_background("beacon_worker", beacon_future);
     }
 
     let app_state = LocalAppState {

crates/mysql/src/connection.rs

@@ -523,7 +523,7 @@ impl<RT: Runtime> Drop for ConvexMySqlPool<RT> {
             return;
         };
         let pool = self.pool.clone();
-        runtime.spawn("mysql_pool_disconnect", async move {
+        runtime.spawn_background("mysql_pool_disconnect", async move {
             let _ = pool.disconnect().await;
             tracing::info!("ConvexMySqlPool pool successfully closed");
         });

crates/runtime/src/prod.rs

@@ -88,6 +88,16 @@ impl SpawnHandle for ThreadHandle {
             Poll::Ready(Ok(()))
         }
     }
+
+    fn detach(mut self: Box<Self>) {
+        self.cancel.take();
+    }
+}
+
+impl Drop for ThreadHandle {
+    fn drop(&mut self) {
+        self.shutdown();
+    }
 }
 
 impl ThreadHandle {