When the JoinHandle is dropped, JOIN_WAKER will never wake up, is this as expected?

[hidva]

use tokio::time::{sleep, Duration};
use tokio::task::JoinHandle;
use std::future::Future;
use std::task::{Context, Poll};
use std::pin::Pin;

// Wait for an asynchronous task to end
struct Wait {
    handle: Option<JoinHandle<()>>,
}

impl Future for Wait {
    type Output = ();
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let muts = self.get_mut();
        if let Some(mut handle) = muts.handle.take() {
            let pinhandle = Pin::new(&mut handle);
            let ret = pinhandle.poll(cx);
            if ret.is_ready() {
                println!("async task is done");
                return Poll::Ready(());
            }
            println!("async task is running");
            // JoinHandle is dropped here.
            return Poll::Pending;
        } else {
            println!("async task is done");
            return Poll::Ready(());
        }
    }
}

#[tokio::main]
async fn main() {
    let handle = tokio::spawn(async {
        println!("before sleep");
        sleep(Duration::from_millis(1000 * 3)).await;
        println!("after sleep");
    });
    let wait = Wait {
        handle: Some(handle),
    };
    wait.await;   // BLOCKING forever!
    println!("wait end");  
}

When I apply the following patch, there is no problem. So I want to know if this is the expected behavior, and if not, I can help fix it.

diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs
index 261dccea..c51160dc 100644
--- a/tokio/src/runtime/task/harness.rs
+++ b/tokio/src/runtime/task/harness.rs
@@ -296,15 +296,15 @@ where
         // We catch panics here in case dropping the future or waking the
         // JoinHandle panics.
         let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
-            if !snapshot.is_join_interested() {
+            if snapshot.has_join_waker() {
+                // Notify the join handle. The previous transition obtains the
+                // lock on the waker cell.
+                self.trailer().wake_join();
+            } else if !snapshot.is_join_interested() {
                 // The `JoinHandle` is not interested in the output of
                 // this task. It is our responsibility to drop the
                 // output.
                 self.core().stage.drop_future_or_output();
-            } else if snapshot.has_join_waker() {
-                // Notify the join handle. The previous transition obtains the
-                // lock on the waker cell.
-                self.trailer().wake_join();
             }
         }));

[Darksonn]

Yes, this is as expected. Pretty much no futures in Tokio will still receive wakeups if dropped.