Trigger CancellationTokens on drop

We'd previously done this for one specific usage via the
`CancelOnDrop` wrapper. This works for the more complicated case where
we need to continue to clone the token. Simpler cases can use the
`DropGuard` directly.

Author: shepmaster

compiler/base/orchestrator/src/worker.rs

@@ -46,7 +46,7 @@ use tokio::{
     sync::mpsc,
     task::JoinSet,
 };
-use tokio_util::sync::CancellationToken;
+use tokio_util::sync::{CancellationToken, DropGuard};
 
 use crate::{
     bincode_input_closed,
@@ -405,7 +405,7 @@ struct ProcessState {
     processes: JoinSet<Result<(), ProcessError>>,
     stdin_senders: HashMap<JobId, mpsc::Sender<String>>,
     stdin_shutdown_tx: mpsc::Sender<JobId>,
-    kill_tokens: HashMap<JobId, CancellationToken>,
+    kill_tokens: HashMap<JobId, DropGuard>,
 }
 
 impl ProcessState {
@@ -458,7 +458,7 @@ impl ProcessState {
 
         let task_set = stream_stdio(worker_msg_tx.clone(), stdin_rx, stdin, stdout, stderr);
 
-        self.kill_tokens.insert(job_id, token.clone());
+        self.kill_tokens.insert(job_id, token.clone().drop_guard());
 
         self.processes.spawn({
             let stdin_shutdown_tx = self.stdin_shutdown_tx.clone();
@@ -510,8 +510,8 @@ impl ProcessState {
     }
 
     fn kill(&mut self, job_id: JobId) {
-        if let Some(token) = self.kill_tokens.get(&job_id) {
-            token.cancel();
+        if let Some(token) = self.kill_tokens.remove(&job_id) {
+            drop(token);
         }
     }
 }

ui/src/server_axum/websocket.rs

@@ -29,7 +29,7 @@ use tokio::{
     task::{AbortHandle, JoinSet},
     time,
 };
-use tokio_util::sync::CancellationToken;
+use tokio_util::sync::{CancellationToken, DropGuard};
 use tracing::{error, info, instrument, warn, Instrument};
 
 #[derive(Debug, serde::Deserialize, serde::Serialize)]
@@ -525,7 +525,7 @@ async fn handle_idle(manager: &mut CoordinatorManager, tx: &ResponseTx) -> Contr
     ControlFlow::Continue(())
 }
 
-type ActiveExecutionInfo = (CancellationToken, Option<mpsc::Sender<String>>);
+type ActiveExecutionInfo = (DropGuard, Option<mpsc::Sender<String>>);
 
 async fn handle_msg(
     txt: &str,
@@ -545,7 +545,10 @@ async fn handle_msg(
 
             let guard = db.clone().start_with_guard("ws.Execute", txt).await;
 
-            active_executions.insert(meta.sequence_number, (token.clone(), Some(execution_tx)));
+            active_executions.insert(
+                meta.sequence_number,
+                (token.clone().drop_guard(), Some(execution_tx)),
+            );
 
             // TODO: Should a single execute / build / etc. session have a timeout of some kind?
             let spawned = manager
@@ -602,11 +605,11 @@ async fn handle_msg(
         }
 
         Ok(ExecuteKill { meta }) => {
-            let Some((token, _)) = active_executions.get(&meta.sequence_number) else {
+            let Some((token, _)) = active_executions.remove(&meta.sequence_number) else {
                 warn!("Received kill for an execution that is no longer active");
                 return;
             };
-            token.cancel();
+            drop(token);
         }
 
         Err(e) => {