Make handling of uni and bidi streams equal

Uni response streams are now accepted as part of the
request lifecycle, which makes it easier to gather stats.

perf/src/client.rs

@@ -84,7 +84,7 @@ async fn run(opt: Opt) -> Result<()> {
 
     let quinn::NewConnection {
         connection,
-        mut uni_streams,
+        uni_streams,
         ..
     } = endpoint
         .connect_with(cfg, &addr, &host_name)?
@@ -93,20 +93,23 @@ async fn run(opt: Opt) -> Result<()> {
 
     info!("established");
 
-    tokio::spawn(async move {
-        while let Some(Ok(stream)) = uni_streams.next().await {
-            tokio::spawn(async move {
-                if let Err(e) = drain_stream(stream).await {
-                    error!("reading response failed: {:#}", e);
-                }
-            });
-        }
-    });
+    let acceptor = UniAcceptor(Arc::new(tokio::sync::Mutex::new(uni_streams)));
 
     let drive_fut = async {
         tokio::try_join!(
-            drive_uni(connection.clone(), opt.uni_requests, opt.upload_size, opt.download_size),
-            drive_bi(connection.clone(), opt.bi_requests, opt.upload_size, opt.download_size)
+            drive_uni(
+                connection.clone(),
+                acceptor,
+                opt.uni_requests,
+                opt.upload_size,
+                opt.download_size
+            ),
+            drive_bi(
+                connection.clone(),
+                opt.bi_requests,
+                opt.upload_size,
+                opt.download_size
+            )
         )
     };
 
@@ -144,6 +147,7 @@ async fn drain_stream(mut stream: quinn::RecvStream) -> Result<()> {
 
 async fn drive_uni(
     connection: quinn::Connection,
+    acceptor: UniAcceptor,
     concurrency: u64,
     upload: u64,
     download: u64,
@@ -153,9 +157,10 @@ async fn drive_uni(
     loop {
         let permit = sem.clone().acquire_owned().await.unwrap();
         let send = connection.open_uni().await?;
+        let acceptor = acceptor.clone();
         trace!("sending request on {}", send.id());
         tokio::spawn(async move {
-            if let Err(e) = request(send, upload, download).await {
+            if let Err(e) = request_uni(send, acceptor, upload, download).await {
                 error!("sending request failed: {:#}", e);
             }
             drop(permit);
@@ -163,6 +168,24 @@ async fn drive_uni(
     }
 }
 
+async fn request_uni(
+    send: quinn::SendStream,
+    acceptor: UniAcceptor,
+    upload: u64,
+    download: u64,
+) -> Result<()> {
+    request(send, upload, download).await?;
+    let recv = {
+        let mut guard = acceptor.0.lock().await;
+        guard
+            .next()
+            .await
+            .ok_or_else(|| anyhow::anyhow!("End of stream"))
+    }??;
+    drain_stream(recv).await?;
+    Ok(())
+}
+
 async fn request(mut send: quinn::SendStream, mut upload: u64, download: u64) -> Result<()> {
     send.write_all(&download.to_be_bytes()).await?;
     let buf = [42; 4 * 1024];
@@ -210,6 +233,9 @@ async fn request_bi(
     Ok(())
 }
 
+#[derive(Clone)]
+struct UniAcceptor(Arc<tokio::sync::Mutex<quinn::IncomingUniStreams>>);
+
 struct SkipServerVerification;
 
 impl SkipServerVerification {