sync: fix mem leak in oneshot on task migration (#1648)

When polling the task, the current waker is saved to the oneshot state.
When the handle is migrated to a new task and polled again, the waker
must be swaped from the old waker to the new waker. In some cases, there
is a potential for the old waker to leak.

This bug was caught by loom with the recently added memory leak
detection.

tokio-sync/src/oneshot.rs

@@ -171,6 +171,8 @@ impl<T> Sender<T> {
                 state = State::unset_tx_task(&inner.state);
 
                 if state.is_closed() {
+                    // Set the flag again so that the waker is released in drop
+                    State::set_tx_task(&inner.state);
                     return Ready(());
                 } else {
                     unsafe { inner.drop_tx_task() };
@@ -359,6 +361,9 @@ impl<T> Inner<T> {
                     // Unset the task
                     state = State::unset_rx_task(&self.state);
                     if state.is_complete() {
+                        // Set the flag again so that the waker is released in drop
+                        State::set_rx_task(&self.state);
+
                         return match unsafe { self.consume_value() } {
                             Some(value) => Ready(Ok(value)),
                             None => Ready(Err(RecvError(()))),