h2-tests: Add testing tools to verify wake calls (#794)

This commit adds wrappers around futures::future helpers and augments
TestFuture to ensure that the underlying futures are notified before
they are polled. This helps to catch bugs where there are missing notify
calls or bad handling of the waker.

The commit then extends the tests to use these helpers instead of the
library functions from futures.

It also ammends the client_requests::recv_too_big_headers test to no
longer use the tokio spawned tasks that were added in #791.

tests/h2-support/src/future_ext.rs

@@ -1,7 +1,9 @@
-use futures::FutureExt;
+use futures::{FutureExt, TryFuture};
 use std::future::Future;
 use std::pin::Pin;
-use std::task::{Context, Poll};
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+use std::task::{Context, Poll, Wake, Waker};
 
 /// Future extension helpers that are useful for tests
 pub trait TestFuture: Future {
@@ -15,9 +17,140 @@ pub trait TestFuture: Future {
     {
         Drive {
             driver: self,
-            future: Box::pin(other),
+            future: other.wakened(),
         }
     }
+
+    fn wakened(self) -> Wakened<Self>
+    where
+        Self: Sized,
+    {
+        Wakened {
+            future: Box::pin(self),
+            woken: Arc::new(AtomicBool::new(true)),
+        }
+    }
+}
+
+/// Wraps futures::future::join to ensure that the futures are only polled if they are woken.
+pub fn join<Fut1, Fut2>(
+    future1: Fut1,
+    future2: Fut2,
+) -> futures::future::Join<Wakened<Fut1>, Wakened<Fut2>>
+where
+    Fut1: Future,
+    Fut2: Future,
+{
+    futures::future::join(future1.wakened(), future2.wakened())
+}
+
+/// Wraps futures::future::join3 to ensure that the futures are only polled if they are woken.
+pub fn join3<Fut1, Fut2, Fut3>(
+    future1: Fut1,
+    future2: Fut2,
+    future3: Fut3,
+) -> futures::future::Join3<Wakened<Fut1>, Wakened<Fut2>, Wakened<Fut3>>
+where
+    Fut1: Future,
+    Fut2: Future,
+    Fut3: Future,
+{
+    futures::future::join3(future1.wakened(), future2.wakened(), future3.wakened())
+}
+
+/// Wraps futures::future::join4 to ensure that the futures are only polled if they are woken.
+pub fn join4<Fut1, Fut2, Fut3, Fut4>(
+    future1: Fut1,
+    future2: Fut2,
+    future3: Fut3,
+    future4: Fut4,
+) -> futures::future::Join4<Wakened<Fut1>, Wakened<Fut2>, Wakened<Fut3>, Wakened<Fut4>>
+where
+    Fut1: Future,
+    Fut2: Future,
+    Fut3: Future,
+    Fut4: Future,
+{
+    futures::future::join4(
+        future1.wakened(),
+        future2.wakened(),
+        future3.wakened(),
+        future4.wakened(),
+    )
+}
+
+/// Wraps futures::future::try_join to ensure that the futures are only polled if they are woken.
+pub fn try_join<Fut1, Fut2>(
+    future1: Fut1,
+    future2: Fut2,
+) -> futures::future::TryJoin<Wakened<Fut1>, Wakened<Fut2>>
+where
+    Fut1: futures::future::TryFuture + Future,
+    Fut2: Future,
+    Wakened<Fut1>: futures::future::TryFuture,
+    Wakened<Fut2>: futures::future::TryFuture<Error = <Wakened<Fut1> as TryFuture>::Error>,
+{
+    futures::future::try_join(future1.wakened(), future2.wakened())
+}
+
+/// Wraps futures::future::select to ensure that the futures are only polled if they are woken.
+pub fn select<A, B>(future1: A, future2: B) -> futures::future::Select<Wakened<A>, Wakened<B>>
+where
+    A: Future + Unpin,
+    B: Future + Unpin,
+{
+    futures::future::select(future1.wakened(), future2.wakened())
+}
+
+/// Wraps futures::future::join_all to ensure that the futures are only polled if they are woken.
+pub fn join_all<I>(iter: I) -> futures::future::JoinAll<Wakened<I::Item>>
+where
+    I: IntoIterator,
+    I::Item: Future,
+{
+    futures::future::join_all(iter.into_iter().map(|f| f.wakened()))
+}
+
+/// A future that only polls the inner future if it has been woken (after the initial poll).
+pub struct Wakened<T> {
+    future: Pin<Box<T>>,
+    woken: Arc<AtomicBool>,
+}
+
+/// A future that only polls the inner future if it has been woken (after the initial poll).
+impl<T> Future for Wakened<T>
+where
+    T: Future,
+{
+    type Output = T::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.get_mut();
+        if !this.woken.load(std::sync::atomic::Ordering::SeqCst) {
+            return Poll::Pending;
+        }
+        this.woken.store(false, std::sync::atomic::Ordering::SeqCst);
+        let my_waker = IfWokenWaker {
+            inner: cx.waker().clone(),
+            wakened: this.woken.clone(),
+        };
+        let my_waker = Arc::new(my_waker).into();
+        let mut cx = Context::from_waker(&my_waker);
+        this.future.as_mut().poll(&mut cx)
+    }
+}
+
+impl Wake for IfWokenWaker {
+    fn wake(self: Arc<Self>) {
+        self.wakened
+            .store(true, std::sync::atomic::Ordering::SeqCst);
+        self.inner.wake_by_ref();
+    }
+}
+
+struct IfWokenWaker {
+    inner: Waker,
+    wakened: Arc<AtomicBool>,
 }
 
 impl<T: Future> TestFuture for T {}
@@ -29,7 +162,7 @@ impl<T: Future> TestFuture for T {}
 /// This is useful for H2 futures that also require the connection to be polled.
 pub struct Drive<'a, T, U> {
     driver: &'a mut T,
-    future: Pin<Box<U>>,
+    future: Wakened<U>,
 }
 
 impl<'a, T, U> Future for Drive<'a, T, U>

tests/h2-support/src/prelude.rs

@@ -35,7 +35,7 @@ pub use {bytes, futures, http, tokio::io as tokio_io, tracing, tracing_subscribe
 pub use futures::{Future, Sink, Stream};
 
 // And our Future extensions
-pub use super::future_ext::TestFuture;
+pub use super::future_ext::{join, join3, join4, join_all, select, try_join, TestFuture};
 
 // Our client_ext helpers
 pub use super::client_ext::SendRequestExt;

tests/h2-tests/tests/client_request.rs

@@ -1,4 +1,4 @@
-use futures::future::{join, join_all, ready, select, Either};
+use futures::future::{ready, Either};
 use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use h2_support::prelude::*;
@@ -849,7 +849,7 @@ async fn recv_too_big_headers() {
     };
 
     let client = async move {
-        let (mut client, conn) = client::Builder::new()
+        let (mut client, mut conn) = client::Builder::new()
             .max_header_list_size(10)
             .handshake::<_, Bytes>(io)
             .await
@@ -863,30 +863,23 @@ async fn recv_too_big_headers() {
         let req1 = client.send_request(request, true);
         // Spawn tasks to ensure that the error wakes up tasks that are blocked
         // waiting for a response.
-        let req1 = tokio::spawn(async move {
+        let req1 = async move {
             let err = req1.expect("send_request").0.await.expect_err("response1");
             assert_eq!(err.reason(), Some(Reason::PROTOCOL_ERROR));
-        });
+        };
 
         let request = Request::builder()
             .uri("https://http2.akamai.com/")
             .body(())
             .unwrap();
 
         let req2 = client.send_request(request, true);
-        let req2 = tokio::spawn(async move {
+        let req2 = async move {
             let err = req2.expect("send_request").0.await.expect_err("response2");
             assert_eq!(err.reason(), Some(Reason::PROTOCOL_ERROR));
-        });
+        };
 
-        let conn = tokio::spawn(async move {
-            conn.await.expect("client");
-        });
-        for err in join_all([req1, req2, conn]).await {
-            if let Some(err) = err.err().and_then(|err| err.try_into_panic().ok()) {
-                std::panic::resume_unwind(err);
-            }
-        }
+        conn.drive(join(req1, req2)).await;
     };
 
     join(srv, client).await;

tests/h2-tests/tests/codec_read.rs

@@ -1,4 +1,3 @@
-use futures::future::join;
 use h2_support::prelude::*;
 
 #[tokio::test]

tests/h2-tests/tests/codec_write.rs

@@ -1,4 +1,3 @@
-use futures::future::join;
 use h2_support::prelude::*;
 
 #[tokio::test]

tests/h2-tests/tests/flow_control.rs

@@ -1,4 +1,3 @@
-use futures::future::{join, join4};
 use futures::{StreamExt, TryStreamExt};
 use h2_support::prelude::*;
 use h2_support::util::yield_once;

tests/h2-tests/tests/ping_pong.rs

@@ -1,5 +1,4 @@
 use futures::channel::oneshot;
-use futures::future::join;
 use futures::StreamExt;
 use h2_support::assert_ping;
 use h2_support::prelude::*;

tests/h2-tests/tests/prioritization.rs

@@ -1,4 +1,3 @@
-use futures::future::{join, select};
 use futures::{pin_mut, FutureExt, StreamExt};
 
 use h2_support::prelude::*;

tests/h2-tests/tests/push_promise.rs

@@ -1,6 +1,4 @@
-use std::iter::FromIterator;
-
-use futures::{future::join, FutureExt as _, StreamExt, TryStreamExt};
+use futures::{StreamExt, TryStreamExt};
 use h2_support::prelude::*;
 
 #[tokio::test]
@@ -52,15 +50,9 @@ async fn recv_push_works() {
             let ps: Vec<_> = p.collect().await;
             assert_eq!(1, ps.len())
         };
-        // Use a FuturesUnordered to poll both tasks but only poll them
-        // if they have been notified.
-        let tasks = futures::stream::FuturesUnordered::from_iter([
-            check_resp_status.boxed(),
-            check_pushed_response.boxed(),
-        ])
-        .collect::<()>();
-
-        h2.drive(tasks).await;
+
+        h2.drive(join(check_resp_status, check_pushed_response))
+            .await;
     };
 
     join(mock, h2).await;

tests/h2-tests/tests/server.rs

@@ -1,6 +1,5 @@
 #![deny(warnings)]
 
-use futures::future::join;
 use futures::StreamExt;
 use h2_support::prelude::*;
 use tokio::io::AsyncWriteExt;

tests/h2-tests/tests/stream_states.rs

@@ -1,6 +1,6 @@
 #![deny(warnings)]
 
-use futures::future::{join, join3, lazy, try_join};
+use futures::future::lazy;
 use futures::{FutureExt, StreamExt, TryStreamExt};
 use h2_support::prelude::*;
 use h2_support::util::yield_once;