util: Fix call_all hang when stream is pending

tower/src/util/call_all/common.rs

@@ -97,20 +97,18 @@ where
                 .expect("Using CallAll after extracing inner Service");
             ready!(svc.poll_ready(cx))?;
 
-            // If it is, gather the next request (if there is one)
-            match this.stream.as_mut().poll_next(cx) {
-                Poll::Ready(r) => match r {
-                    Some(req) => {
-                        this.queue.push(svc.call(req));
-                    }
-                    None => {
-                        // We're all done once any outstanding requests have completed
-                        *this.eof = true;
-                    }
-                },
-                Poll::Pending => {
-                    // TODO: We probably want to "release" the slot we reserved in Svc here.
-                    // It may be a while until we get around to actually using it.
+            // If it is, gather the next request (if there is one), or return `Pending` if the
+            // stream is not ready.
+            // TODO: We probably want to "release" the slot we reserved in Svc if the
+            // stream returns `Pending`. It may be a while until we get around to actually
+            // using it.
+            match ready!(this.stream.as_mut().poll_next(cx)) {
+                Some(req) => {
+                    this.queue.push(svc.call(req));
+                }
+                None => {
+                    // We're all done once any outstanding requests have completed
+                    *this.eof = true;
                 }
             }
         }

tower/tests/util/call_all.rs

@@ -143,3 +143,24 @@ async fn unordered() {
         .unwrap();
     assert!(v.is_none());
 }
+
+#[tokio::test]
+async fn pending() {
+    let _t = support::trace_init();
+
+    let (mock, mut handle) = mock::pair::<_, &'static str>();
+
+    let mut task = task::spawn(());
+
+    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+    let ca = mock.call_all(support::IntoStream::new(rx));
+    pin_mut!(ca);
+
+    assert_pending!(task.enter(|cx, _| ca.as_mut().poll_next(cx)));
+    tx.send("req").unwrap();
+    assert_pending!(task.enter(|cx, _| ca.as_mut().poll_next(cx)));
+    assert_request_eq!(handle, "req").send_response("res");
+    let res = assert_ready!(task.enter(|cx, _| ca.as_mut().poll_next(cx)));
+    assert_eq!(res.transpose().unwrap(), Some("res"));
+    assert_pending!(task.enter(|cx, _| ca.as_mut().poll_next(cx)));
+}