Abort plain Tokio tasks on drop

`JoinSet`s already abort all their tasks when the set is dropped, so
nothing to worry about there.

Blocking tasks aren't really abortable, so I didn't change those.

Author: shepmaster

compiler/base/orchestrator/Cargo.lock

@@ -20,7 +20,7 @@ snafu = { version = "0.8.0", default-features = false, features = ["futures", "s
 strum_macros = { version = "0.26.1", default-features = false }
 tokio = { version = "1.28", default-features = false, features = ["fs", "io-std", "io-util", "macros", "process", "rt", "time", "sync"] }
 tokio-stream = { version = "0.1.14", default-features = false }
-tokio-util = { version = "0.7.8", default-features = false, features = ["io", "io-util"] }
+tokio-util = { version = "0.7.8", default-features = false, features = ["io", "io-util", "rt"] }
 toml = { version = "0.8.2", default-features = false, features = ["parse", "display"] }
 tracing = { version = "0.1.37", default-features = false, features = ["attributes"] }
 

compiler/base/orchestrator/Cargo.toml

@@ -15,12 +15,12 @@ use tokio::{
     process::{Child, ChildStdin, ChildStdout, Command},
     select,
     sync::{mpsc, oneshot, OnceCell},
-    task::{JoinHandle, JoinSet},
+    task::JoinSet,
     time::{self, MissedTickBehavior},
     try_join,
 };
 use tokio_stream::wrappers::ReceiverStream;
-use tokio_util::{io::SyncIoBridge, sync::CancellationToken};
+use tokio_util::{io::SyncIoBridge, sync::CancellationToken, task::AbortOnDropHandle};
 use tracing::{error, info, info_span, instrument, trace, trace_span, warn, Instrument};
 
 use crate::{
@@ -30,7 +30,7 @@ use crate::{
         ExecuteCommandResponse, JobId, Multiplexed, OneToOneResponse, ReadFileRequest,
         ReadFileResponse, SerializedError2, WorkerMessage, WriteFileRequest,
     },
-    DropErrorDetailsExt,
+    DropErrorDetailsExt, TaskAbortExt as _,
 };
 
 pub mod limits;
@@ -1161,7 +1161,7 @@ impl Drop for CancelOnDrop {
 #[derive(Debug)]
 struct Container {
     permit: Box<dyn ContainerPermit>,
-    task: JoinHandle<Result<()>>,
+    task: AbortOnDropHandle<Result<()>>,
     kill_child: TerminateContainer,
     modify_cargo_toml: ModifyCargoToml,
     commander: Commander,
@@ -1186,7 +1186,8 @@ impl Container {
 
         let (command_tx, command_rx) = mpsc::channel(8);
         let demultiplex_task =
-            tokio::spawn(Commander::demultiplex(command_rx, from_worker_rx).in_current_span());
+            tokio::spawn(Commander::demultiplex(command_rx, from_worker_rx).in_current_span())
+                .abort_on_drop();
 
         let task = tokio::spawn(
             async move {
@@ -1216,7 +1217,8 @@ impl Container {
                 Ok(())
             }
             .in_current_span(),
-        );
+        )
+        .abort_on_drop();
 
         let commander = Commander {
             to_worker_tx,
@@ -1865,7 +1867,8 @@ impl Container {
                 }
             }
             .instrument(trace_span!("cargo task").or_current())
-        });
+        })
+        .abort_on_drop();
 
         Ok(SpawnCargo {
             permit,
@@ -2128,7 +2131,7 @@ pub enum DoRequestError {
 
 struct SpawnCargo {
     permit: Box<dyn ProcessPermit>,
-    task: JoinHandle<Result<ExecuteCommandResponse, SpawnCargoError>>,
+    task: AbortOnDropHandle<Result<ExecuteCommandResponse, SpawnCargoError>>,
     stdin_tx: mpsc::Sender<String>,
     stdout_rx: mpsc::Receiver<String>,
     stderr_rx: mpsc::Receiver<String>,

compiler/base/orchestrator/src/coordinator.rs

@@ -4,6 +4,16 @@ pub mod coordinator;
 mod message;
 pub mod worker;
 
+pub trait TaskAbortExt<T>: Sized {
+    fn abort_on_drop(self) -> tokio_util::task::AbortOnDropHandle<T>;
+}
+
+impl<T> TaskAbortExt<T> for tokio::task::JoinHandle<T> {
+    fn abort_on_drop(self) -> tokio_util::task::AbortOnDropHandle<T> {
+        tokio_util::task::AbortOnDropHandle::new(self)
+    }
+}
+
 pub trait DropErrorDetailsExt<T> {
     fn drop_error_details(self) -> Result<T, tokio::sync::mpsc::error::SendError<()>>;
 }

compiler/base/orchestrator/src/lib.rs

@@ -55,7 +55,7 @@ use crate::{
         ExecuteCommandResponse, JobId, Multiplexed, ReadFileRequest, ReadFileResponse,
         SerializedError2, WorkerMessage, WriteFileRequest, WriteFileResponse,
     },
-    DropErrorDetailsExt,
+    DropErrorDetailsExt as _, TaskAbortExt as _,
 };
 
 pub async fn listen(project_dir: impl Into<PathBuf>) -> Result<(), Error> {
@@ -66,14 +66,16 @@ pub async fn listen(project_dir: impl Into<PathBuf>) -> Result<(), Error> {
     let mut io_tasks = spawn_io_queue(coordinator_msg_tx, worker_msg_rx);
 
     let (process_tx, process_rx) = mpsc::channel(8);
-    let process_task = tokio::spawn(manage_processes(process_rx, project_dir.clone()));
+    let process_task =
+        tokio::spawn(manage_processes(process_rx, project_dir.clone())).abort_on_drop();
 
     let handler_task = tokio::spawn(handle_coordinator_message(
         coordinator_msg_rx,
         worker_msg_tx,
         project_dir,
         process_tx,
-    ));
+    ))
+    .abort_on_drop();
 
     select! {
         Some(io_task) = io_tasks.join_next() => {

compiler/base/orchestrator/src/worker.rs

@@ -2,7 +2,7 @@ use futures::{
     future::{Fuse, FusedFuture as _},
     FutureExt as _,
 };
-use orchestrator::DropErrorDetailsExt as _;
+use orchestrator::{DropErrorDetailsExt as _, TaskAbortExt as _};
 use snafu::prelude::*;
 use std::{
     future::Future,
@@ -13,9 +13,9 @@ use std::{
 use tokio::{
     select,
     sync::{mpsc, oneshot},
-    task::JoinHandle,
     time,
 };
+use tokio_util::task::AbortOnDropHandle;
 use tracing::warn;
 
 const ONE_HUNDRED_MILLISECONDS: Duration = Duration::from_millis(100);
@@ -48,12 +48,12 @@ where
 {
     pub fn spawn<Fut>(
         f: impl FnOnce(mpsc::Receiver<CacheTaskItem<T, E>>) -> Fut,
-    ) -> (JoinHandle<()>, Self)
+    ) -> (AbortOnDropHandle<()>, Self)
     where
         Fut: Future<Output = ()> + Send + 'static,
     {
         let (tx, rx) = mpsc::channel(8);
-        let task = tokio::spawn(f(rx));
+        let task = tokio::spawn(f(rx)).abort_on_drop();
         let cache_tx = CacheTx(tx);
         (task, cache_tx)
     }
@@ -148,7 +148,8 @@ where
                         let new_value = generator().await.map_err(CacheError::from);
                         CacheInfo::build(new_value)
                     }
-                });
+                })
+                .abort_on_drop();
 
                 new_value.set(new_value_task.fuse());
             }