From 35cf37c7823fbda70dcd9ea34d42adb0ab8af1fc Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Sun, 27 Aug 2023 17:16:18 -0500
Subject: [PATCH] Propagate jitterbug cancel to bachata

---
 examples/time.rs |  4 ++++
 src/runtime.rs   | 19 +++++++++++++++++--
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/examples/time.rs b/examples/time.rs
index 4a44ea1..8b9957b 100644
--- a/examples/time.rs
+++ b/examples/time.rs
@@ -20,6 +20,8 @@ fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
                         handle.await?;
                     }
 
+                    println!("{i} joined");
+
                     Ok(()) as Result<_, jive::task::sync::JoinError>
                 })
             })
@@ -29,6 +31,8 @@ fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
             handle.await??;
         }
 
+        println!("all joined");
+
         Ok(())
     })
 }
diff --git a/src/runtime.rs b/src/runtime.rs
index 0760fcc..4646158 100644
--- a/src/runtime.rs
+++ b/src/runtime.rs
@@ -92,8 +92,23 @@ impl RuntimeHandle {
         future: impl Future<Output = T> + 'static,
     ) -> jitterbug::JoinHandle<T> {
         let (tx, rx) = jitterbug::oneshot();
-        bachata::spawn(async move { tx.send(future.await) });
-        self.spawn(async move { rx.await.unwrap() })
+        let (cancel_tx, cancel_rx) = jitterbug::oneshot::<()>();
+
+        bachata::spawn(async move {
+            let future = std::pin::pin!(future);
+
+            match select::select(cancel_rx, future).await {
+                select::Either::Left(_) => (), // canceled
+                select::Either::Right(item) => {
+                    let _ = tx.send(item);
+                }
+            }
+        });
+        self.spawn(async move {
+            let out = rx.await.unwrap();
+            drop(cancel_tx);
+            out
+        })
     }
 
     pub fn spawn_blocking<T: Send + 'static>(