Repository: tokio-rs/tokio
File: tokio/src/runtime/task/state.rs (Rust)

Code before:
        self.fetch_update_action(|mut snapshot| {
            assert!(snapshot.is_join_interested());

            let mut transition = TransitionToJoinHandleDrop {
                drop_waker: false,
                drop_output: false,
            };

            snapshot.unset_join_interested();

            if !snapshot.is_complete() {
                // If `COMPLETE` is unset we also unset `JOIN_WAKER` to give the
                // `JoinHandle` exclusive access to the waker following rule 6 in task/mod.rs.
                // The `JoinHandle` will drop the waker if it has exclusive access
                // to drop it.
                snapshot.unset_join_waker();
            } else {
                // If `COMPLETE` is set the task is completed so the `JoinHandle` is responsible
                // for dropping the output.
                transition.drop_output = true;
            }

            if !snapshot.is_join_waker_set() {
                // If the `JOIN_WAKER` bit is unset and the `JOIN_HANDLE` has exclusive access to
                // the the join waker and should drop it following this transition.
                transition.drop_waker = true;
            }

            (transition, Some(snapshot))
        })
    }

    /// Sets the `JOIN_WAKER` bit.
    ///
    /// Returns `Ok` if the bit is set, `Err` otherwise. This operation fails if
    /// the task has completed.
    pub(super) fn set_join_waker(&self) -> UpdateResult {
        self.fetch_update(|curr| {
            assert!(curr.is_join_interested());
            assert!(!curr.is_join_waker_set());

            if curr.is_complete() {
                return None;
            }

            let mut next = curr;
            next.set_join_waker();

            Some(next)
        })
    }

Review comment (suggestion):
It's probably good to mention in the comment that there are two cases where we drop the join waker:

1. We just unset `JOIN_WAKER` since the task was not completed.
2. Or if `COMPLETE` is set, then `JOIN_WAKER` may already have been unset by the runtime.

So this can in fact be true even if `COMPLETE == 1`.

Code after:
        self.fetch_update_action(|mut snapshot| {
            assert!(snapshot.is_join_interested());

            let mut transition = TransitionToJoinHandleDrop {
                drop_waker: false,
                drop_output: false,
            };

            snapshot.unset_join_interested();

            if !snapshot.is_complete() {
                // If `COMPLETE` is unset we also unset `JOIN_WAKER` to give the
                // `JoinHandle` exclusive access to the waker following rule 6 in task/mod.rs.
                // The `JoinHandle` will drop the waker if it has exclusive access
                // to drop it.
                snapshot.unset_join_waker();
            } else {
                // If `COMPLETE` is set the task is completed so the `JoinHandle` is responsible
                // for dropping the output.
                transition.drop_output = true;
            }

            if !snapshot.is_join_waker_set() {
                // If the `JOIN_WAKER` bit is unset and the `JOIN_HANDLE` has exclusive access to
                // the join waker and should drop it following this transition.
                // This might happen in two situations:
                //  1. The task is not completed and we just unset the `JOIN_WAKer` above in this
                //     function.
                //  2. The task is completed. In that case the `JOIN_WAKER` bit was already unset
                //     by the runtime during completion.
                transition.drop_waker = true;
            }

            (transition, Some(snapshot))
        })
    }

    /// Sets the `JOIN_WAKER` bit.
    ///
    /// Returns `Ok` if the bit is set, `Err` otherwise. This operation fails if
    /// the task has completed.
    pub(super) fn set_join_waker(&self) -> UpdateResult {
        self.fetch_update(|curr| {
            assert!(curr.is_join_interested());
            assert!(!curr.is_join_waker_set());

            if curr.is_complete() {
                return None;
            }

            let mut next = curr;