From 2eed6d00f551c074f48054a1412fa85e9fbe13f7 Mon Sep 17 00:00:00 2001
From: Jon Gjengset <jon@thesquareplanet.com>
Date: Wed, 12 Feb 2020 15:53:56 -0500
Subject: [PATCH] Avoid race in tcp_accept::no_extra_poll (#2236)

---
 tokio/tests/tcp_accept.rs | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tokio/tests/tcp_accept.rs b/tokio/tests/tcp_accept.rs
index a5a0e2d56..ff62fb96a 100644
--- a/tokio/tests/tcp_accept.rs
+++ b/tokio/tests/tcp_accept.rs
@@ -2,7 +2,7 @@
 #![cfg(feature = "full")]
 
 use tokio::net::{TcpListener, TcpStream};
-use tokio::sync::oneshot;
+use tokio::sync::{mpsc, oneshot};
 use tokio_test::assert_ok;
 
 use std::net::{IpAddr, SocketAddr};
@@ -72,6 +72,7 @@ async fn no_extra_poll() {
     let addr = listener.local_addr().unwrap();
 
     let (tx, rx) = oneshot::channel();
+    let (accepted_tx, mut accepted_rx) = mpsc::unbounded_channel();
 
     tokio::spawn(async move {
         let mut incoming = TrackPolls {
@@ -79,7 +80,9 @@ async fn no_extra_poll() {
             s: listener.incoming(),
         };
         assert_ok!(tx.send(Arc::clone(&incoming.npolls)));
-        while let Some(_) = incoming.next().await {}
+        while let Some(_) = incoming.next().await {
+            accepted_tx.send(()).unwrap();
+        }
     });
 
     let npolls = assert_ok!(rx.await);
@@ -89,7 +92,7 @@ async fn no_extra_poll() {
     assert_eq!(npolls.load(SeqCst), 1);
 
     let _ = assert_ok!(TcpStream::connect(&addr).await);
-    tokio::task::yield_now().await;
+    accepted_rx.next().await.unwrap();
 
     // should have been polled twice more: once to yield Some(), then once to yield Pending
     assert_eq!(npolls.load(SeqCst), 1 + 2);