1 files changed, 21 insertions, 10 deletions

diff --git a/src/join_handle.rs b/src/join_handle.rs
index f4710cc..10a189e 100644
--- a/src/join_handle.rs
+++ b/src/join_handle.rs
@@ -8,7 +8,6 @@ use core::task::{Context, Poll, Waker};
 
 use crate::header::Header;
 use crate::state::*;
-use crate::utils::abort_on_panic;
 
 /// A handle that awaits the result of a task.
 ///
@@ -199,6 +198,23 @@ impl<R, T> Future for JoinHandle<R, T> {
             loop {
                 // If the task has been closed, notify the awaiter and return `None`.
                 if state & CLOSED != 0 {
+                    // If the task is scheduled or running, we need to wait until its future is
+                    // dropped.
+                    if state & (SCHEDULED | RUNNING) != 0 {
+                        // Replace the waker with one associated with the current task.
+                        (*header).register(cx.waker());
+
+                        // Reload the state after registering. It is possible changes occurred just
+                        // before registration so we need to check for that.
+                        state = (*header).state.load(Ordering::Acquire);
+
+                        // If the task is still scheduled or running, we need to wait because its
+                        // future is not dropped yet.
+                        if state & (SCHEDULED | RUNNING) != 0 {
+                            return Poll::Pending;
+                        }
+                    }
+
                     // Even though the awaiter is most likely the current task, it could also be
                     // another task.
                     (*header).notify(Some(cx.waker()));
@@ -207,21 +223,16 @@ impl<R, T> Future for JoinHandle<R, T> {
 
                 // If the task is not completed, register the current task.
                 if state & COMPLETED == 0 {
-                    // Replace the waker with one associated with the current task. We need a
-                    // safeguard against panics because dropping the previous waker can panic.
-                    abort_on_panic(|| {
-                        (*header).register(cx.waker());
-                    });
+                    // Replace the waker with one associated with the current task.
+                    (*header).register(cx.waker());
 
                     // Reload the state after registering. It is possible that the task became
                     // completed or closed just before registration so we need to check for that.
                     state = (*header).state.load(Ordering::Acquire);
 
-                    // If the task has been closed, return `None`. We do not need to notify the
-                    // awaiter here, since we have replaced the waker above, and the executor can
-                    // only set it back to `None`.
+                    // If the task has been closed, restart.
                     if state & CLOSED != 0 {
-                        return Poll::Ready(None);
+                        continue;
                     }
 
                     // If the task is still not completed, we're blocked on it.