fix: Shared must relinquish control to the executor if repolled

futures-util/src/future/future/shared.rs

@@ -66,9 +66,8 @@ where
 
 const IDLE: usize = 0;
 const POLLING: usize = 1;
-const REPOLL: usize = 2;
-const COMPLETE: usize = 3;
-const POISONED: usize = 4;
+const COMPLETE: usize = 2;
+const POISONED: usize = 3;
 
 const NULL_WAKER_KEY: usize = usize::max_value();
 
@@ -196,7 +195,7 @@ where
             IDLE => {
                 // Lock acquired, fall through
             }
-            POLLING | REPOLL => {
+            POLLING => {
                 // Another task is currently polling, at this point we just want
                 // to ensure that the waker for this task is registered
                 this.inner = Some(inner);
@@ -227,35 +226,27 @@ where
 
         let _reset = Reset(&inner.notifier.state);
 
-        let output = loop {
+        let output = {
             let future = unsafe {
                 match &mut *inner.future_or_output.get() {
                     FutureOrOutput::Future(fut) => Pin::new_unchecked(fut),
                     _ => unreachable!(),
                 }
             };
 
-            let poll = future.poll(&mut cx);
-
-            match poll {
+            match future.poll(&mut cx) {
                 Poll::Pending => {
-                    let state = &inner.notifier.state;
-                    match state.compare_and_swap(POLLING, IDLE, SeqCst) {
+                    match inner.notifier.state.compare_and_swap(POLLING, IDLE, SeqCst) {
                         POLLING => {
                             // Success
                             drop(_reset);
                             this.inner = Some(inner);
                             return Poll::Pending;
                         }
-                        REPOLL => {
-                            // Was woken since: Gotta poll again!
-                            let prev = state.swap(POLLING, SeqCst);
-                            assert_eq!(prev, REPOLL);
-                        }
                         _ => unreachable!(),
                     }
                 }
-                Poll::Ready(output) => break output,
+                Poll::Ready(output) => output,
             }
         };
 
@@ -313,8 +304,6 @@ where
 
 impl ArcWake for Notifier {
     fn wake_by_ref(arc_self: &Arc<Self>) {
-        arc_self.state.compare_and_swap(POLLING, REPOLL, SeqCst);
-
         let wakers = &mut *arc_self.wakers.lock().unwrap();
         if let Some(wakers) = wakers.as_mut() {
             for (_key, opt_waker) in wakers {

futures/tests/shared.rs

@@ -89,7 +89,6 @@ fn drop_on_one_task_ok() {
     t2.join().unwrap();
 }
 
-
 #[cfg(feature = "executor")] // executor::
 #[test]
 fn drop_in_poll() {
@@ -104,7 +103,8 @@ fn drop_in_poll() {
     let future1 = future::lazy(move |_| {
         slot2.replace(None); // Drop future
         1
-    }).shared();
+    })
+    .shared();
 
     let future2 = LocalFutureObj::new(Box::new(future1.clone()));
     slot1.replace(Some(future2));
@@ -141,7 +141,9 @@ fn peek() {
     }
 
     // Once the Shared has been polled, the value is peekable on the clone.
-    spawn.spawn_local_obj(LocalFutureObj::new(Box::new(f1.map(|_| ())))).unwrap();
+    spawn
+        .spawn_local_obj(LocalFutureObj::new(Box::new(f1.map(|_| ()))))
+        .unwrap();
     local_pool.run();
     for _ in 0..2 {
         assert_eq!(*f2.peek().unwrap(), Ok(42));
@@ -191,3 +193,30 @@ fn dont_do_unnecessary_clones_on_output() {
     assert_eq!(block_on(rx.clone()).unwrap().0.get(), 2);
     assert_eq!(block_on(rx).unwrap().0.get(), 2);
 }
+
+#[cfg(all(feature = "alloc", feature = "executor"))] // channel:: + executor::
+#[test]
+fn shared_future_that_wakes_itself_until_pending_is_returned() {
+    use futures::executor::block_on;
+    use futures::future::FutureExt;
+    use std::cell::Cell;
+    use std::task::Poll;
+
+    let proceed = Cell::new(false);
+    let fut = futures::future::poll_fn(|cx| {
+        if proceed.get() {
+            Poll::Ready(())
+        } else {
+            cx.waker().wake_by_ref();
+            Poll::Pending
+        }
+    })
+    .shared();
+
+    // The join future can only complete if the second future gets a chance to run after the first
+    // has returned pending
+    assert_eq!(
+        block_on(futures::future::join(fut, async { proceed.set(true) })),
+        ((), ())
+    );
+}