Tokio runtime stuck when using shared futures

[pranjalssh]

Version

0.2.14 and above

Platform

Darwin, Linux

Description

Tokio runtime gets stuck when using FutureExt::shared.

Short repro: https://gist.github.com/rust-play/20bc6b3c2490cfc721012d3ffb279d81

The code just spawns 2 tasks on runtime. Task Y send queries to task X, task X responds to them and task Y awaits on the replies. Task Y gets stuck almost 90% of the time, when one spawned task completes exactly 128 awaits.

This relates well to #2160 and the new coop budget introduced. And is only reproducible on 0.2.14 and above.

use tokio::sync::mpsc;
use tokio::sync::oneshot;
extern crate rand;
use futures::future::FutureExt;

use rand::Rng;


async fn lmao_test() {
    let (gtx, mut grx): (
        mpsc::UnboundedSender<oneshot::Sender<bool>>,
        mpsc::UnboundedReceiver<oneshot::Sender<bool>>,
    ) = mpsc::unbounded_channel();
    tokio::spawn(async move {
        let mut ctr = 0;
        loop {
            println!("WAITING");
            if let Some(i) = grx.recv().await {
                i.send(false).unwrap();
                ctr += 1;
                println!("SENT {}", ctr);
            } else {
                println!("FINISHED");
                break;
            }
        }
    });
    let mut rng = rand::thread_rng();
    let t = rng.gen_range(200, 500);
    tokio::spawn(async move {
        let mut futs = vec![];
        println!("{} times", t);
        for _ in 0..t {
            let (tx, rx) = oneshot::channel::<bool>();
            gtx.send(tx).unwrap();
            // If you remove shared(), it works
            futs.push(Box::pin(async move { rx.await.unwrap() }).shared());
        }
        let mut ctr = 0;
        for f in futs {
            println!("start {}", ctr);
            if f.await {
                panic!("iwkms");
            }
            println!("end {}", ctr);
            ctr = ctr + 1;
        }
    })
    .await;
}
fn main() {
     let mut rt = tokio::runtime::Builder::new()
        .basic_scheduler()
        .enable_all()
        .build()
        .unwrap();
    rt.block_on(async {lmao_test().await;});
}

[Darksonn]

This is a bug in Shared. The issue is the REPOLL case of the match, which always goes around the loop leading to an infinite loop of polling it again and again and again.

let output = loop {
    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 {
        Poll::Pending => {
            let state = &inner.notifier.state;
            match 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,
    }
};

click to see in repository

[Darksonn]

Closing this issue in favor of futures#2130.