Repository: tokio-rs/tokio
File: tokio/tests/rt_common.rs (Rust)

Code before:
                        std::thread::sleep(std::time::Duration::from_millis(1));
                    }
                });
            }

            barrier.wait();

            let (fail_test, fail_test_recv) = oneshot::channel::<()>();
            let flag_clone = flag.clone();
            let jh = tokio::spawn(async move {
                // Create a TCP listener
                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
                let addr = listener.local_addr().unwrap();

                tokio::join!(
                    async {
                        // Done in a blocking manner intentionally.
                        let _socket = std::net::TcpStream::connect(addr).unwrap();

                        // Yield until connected
                        let mut cnt = 0;
                        while !flag_clone.load(SeqCst){
                            if use_coop {
                                // Consume a good chunk of budget, which should
                                // force at least one yield.
                                for _ in 0..200 {
                                    tokio::task::consume_budget().await;
                                }
                            } else {
                                tokio::task::yield_now().await;
                            }
                            cnt += 1;

                            if cnt >= 10 {
                                // yielded too many times; report failure and
                                // sleep forever so that the `fail_test` branch
                                // of the `select!` below triggers.
                                let _ = fail_test.send(());
                                futures::future::pending::<()>().await;
                                break;
                            }
                        }
                    },
                    async {
                        let _ = listener.accept().await.unwrap();
                        flag_clone.store(true, SeqCst);
                    }
                );
            });

            // Wait until the spawned task completes or fails. If no message is

Review comment (suggestion):
We have a fair number of tests that declare a `const BUDGET: usize = 128;` constant. It would be good to do the same here so we can grep for tests that assume what the budget is.

Code after:

                    while !flag.load(SeqCst) {
                        std::thread::sleep(std::time::Duration::from_millis(1));
                    }
                });
            }

            barrier.wait();

            let (fail_test, fail_test_recv) = oneshot::channel::<()>();
            let flag_clone = flag.clone();
            let jh = tokio::spawn(async move {
                // Create a TCP listener
                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
                let addr = listener.local_addr().unwrap();

                tokio::join!(
                    async {
                        // Done in a blocking manner intentionally.
                        let _socket = std::net::TcpStream::connect(addr).unwrap();

                        // Yield until connected
                        let mut cnt = 0;
                        while !flag_clone.load(SeqCst){
                            if use_coop {
                                // Consume a good chunk of budget, which should
                                // force at least one yield.
                                for _ in 0..BUDGET {
                                    tokio::task::consume_budget().await;
                                }
                            } else {
                                tokio::task::yield_now().await;
                            }
                            cnt += 1;

                            if cnt >= 10 {
                                // yielded too many times; report failure and
                                // sleep forever so that the `fail_test` branch
                                // of the `select!` below triggers.
                                let _ = fail_test.send(());
                                futures::future::pending::<()>().await;
                                break;
                            }
                        }
                    },
                    async {
                        let _ = listener.accept().await.unwrap();
                        flag_clone.store(true, SeqCst);
                    }
                );
            });