From e082634b5e4ce0e89c6107ae61be3ab5b0e95445 Mon Sep 17 00:00:00 2001
From: dignifiedquire <me@dignifiedquire.com>
Date: Fri, 24 Apr 2020 22:53:39 +0200
Subject: [PATCH] fix spawning

---
 Cargo.toml                 |  2 +-
 src/task/block_on.rs       |  6 +++---
 src/task/builder.rs        | 12 +++++++-----
 src/task/join_handle.rs    | 29 +++++++++++++++++++++++------
 src/task/spawn_blocking.rs |  6 +++---
 5 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 45662bd..b850b54 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -56,7 +56,7 @@ alloc = [
 
 [dependencies]
 async-attributes = { version = "1.1.1", optional = true }
-async-task = { version = "1.3.1", optional = true }
+async-task = { version = "3.0.0", optional = true }
 broadcaster = { version = "1.0.0", optional = true }
 crossbeam-channel = { version = "0.4.2", optional = true }
 crossbeam-deque = { version = "0.7.3", optional = true }
diff --git a/src/task/block_on.rs b/src/task/block_on.rs
index 41e0ca0..65d6541 100644
--- a/src/task/block_on.rs
+++ b/src/task/block_on.rs
@@ -42,9 +42,9 @@ where
 
     let wrapped_future = async move {
         // Drop task-locals on exit.
-        // defer! {
-        //     Task::get_current(|t| unsafe { t.drop_locals() });
-        // }
+        defer! {
+            Task::get_current(|t| unsafe { t.drop_locals() });
+        }
 
         // Log completion on exit.
         defer! {
diff --git a/src/task/builder.rs b/src/task/builder.rs
index 898308c..de3e6ac 100644
--- a/src/task/builder.rs
+++ b/src/task/builder.rs
@@ -42,9 +42,9 @@ impl Builder {
 
         let wrapped_future = async move {
             // Drop task-locals on exit.
-            // defer! {
-            //     Task::get_current(|t| unsafe { t.drop_locals() });
-            // }
+            defer! {
+                Task::get_current(|t| unsafe { t.drop_locals() });
+            }
 
             // Log completion on exit.
             defer! {
@@ -57,7 +57,9 @@ impl Builder {
 
         once_cell::sync::Lazy::force(&crate::rt::RUNTIME);
 
-        let task = smol::Task::spawn(wrapped_future);
-        Ok(JoinHandle::new(task))
+        // FIXME: figure out how to set the current task.
+
+        let smol_task = smol::Task::spawn(wrapped_future).detach();
+        Ok(JoinHandle::new(smol_task, task))
     }
 }
diff --git a/src/task/join_handle.rs b/src/task/join_handle.rs
index 72fb2e0..3a63271 100644
--- a/src/task/join_handle.rs
+++ b/src/task/join_handle.rs
@@ -12,12 +12,18 @@ use crate::task::{Context, Poll, Task};
 ///
 /// [spawned]: fn.spawn.html
 #[derive(Debug)]
-pub struct JoinHandle<T>(smol::Task<T>);
+pub struct JoinHandle<T> {
+    handle: Option<async_task::JoinHandle<T, ()>>,
+    task: Task,
+}
 
 impl<T> JoinHandle<T> {
     /// Creates a new `JoinHandle`.
-    pub(crate) fn new(inner: smol::Task<T>) -> JoinHandle<T> {
-        JoinHandle(inner)
+    pub(crate) fn new(inner: async_task::JoinHandle<T, ()>, task: Task) -> JoinHandle<T> {
+        JoinHandle {
+            handle: Some(inner),
+            task,
+        }
     }
 
     /// Returns a handle to the underlying task.
@@ -36,7 +42,14 @@ impl<T> JoinHandle<T> {
     /// #
     /// # })
     pub fn task(&self) -> &Task {
-        todo!()
+        &self.task
+    }
+
+    /// Cancel this task.
+    pub async fn cancel(mut self) -> Option<T> {
+        let handle = self.handle.take().unwrap();
+        handle.cancel();
+        handle.await
     }
 }
 
@@ -44,7 +57,11 @@ impl<T> Future for JoinHandle<T> {
     type Output = T;
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        dbg!("poll joinhandle");
-        Pin::new(&mut self.0).poll(cx)
+        match Pin::new(&mut self.handle.as_mut().unwrap()).poll(cx) {
+            Poll::Pending => Poll::Pending,
+            Poll::Ready(output) => {
+                Poll::Ready(output.expect("cannot await the result of a panicked task"))
+            }
+        }
     }
 }
diff --git a/src/task/spawn_blocking.rs b/src/task/spawn_blocking.rs
index d7b4fd0..054f3fd 100644
--- a/src/task/spawn_blocking.rs
+++ b/src/task/spawn_blocking.rs
@@ -1,4 +1,4 @@
-use crate::task::JoinHandle;
+use crate::task::{JoinHandle, Task};
 
 /// Spawns a blocking task.
 ///
@@ -37,6 +37,6 @@ where
 {
     once_cell::sync::Lazy::force(&crate::rt::RUNTIME);
 
-    let handle = smol::Task::blocking(async move { f() });
-    JoinHandle::new(handle)
+    let handle = smol::Task::blocking(async move { f() }).detach();
+    JoinHandle::new(handle, Task::new(None))
 }