Address flaky execute_kill test

The primary loop in `spawn_cargo_task` works by getting messages from
the worker and then handing them off to the stdout / stderr / status
channels. If one channel is full, then everything will queue behind
it, resulting in all work stopping.

The test was waiting on the stdout channel and had dropped the status
channel, but nothing was processing the stderr channel. This means
that if enough[^1] stderr packets had been delivered before any
stdout, the test would fail.

To reproduce this, you can hack the loop in `spawn_cargo_task` so that
it delivers every `stderr` packet multiple times, forcing the channel
to be filled immediately.

[^1]: Over eight with the current setup.

Author: shepmaster

compiler/base/orchestrator/src/coordinator.rs

@@ -1,7 +1,7 @@
 use futures::{
     future::{BoxFuture, OptionFuture},
     stream::BoxStream,
-    Future, FutureExt, StreamExt,
+    Future, FutureExt, Stream, StreamExt,
 };
 use once_cell::sync::Lazy;
 use serde::Deserialize;
@@ -764,8 +764,24 @@ impl<T> WithOutput<T> {
     where
         F: Future<Output = Result<T, E>>,
     {
-        let stdout = ReceiverStream::new(stdout_rx).collect();
-        let stderr = ReceiverStream::new(stderr_rx).collect();
+        Self::try_absorb_stream(
+            task,
+            ReceiverStream::new(stdout_rx),
+            ReceiverStream::new(stderr_rx),
+        )
+        .await
+    }
+
+    async fn try_absorb_stream<F, E>(
+        task: F,
+        stdout_rx: impl Stream<Item = String>,
+        stderr_rx: impl Stream<Item = String>,
+    ) -> Result<WithOutput<T>, E>
+    where
+        F: Future<Output = Result<T, E>>,
+    {
+        let stdout = stdout_rx.collect();
+        let stderr = stderr_rx.collect();
 
         let (result, stdout, stderr) = join!(task, stdout, stderr);
         let response = result?;
@@ -815,6 +831,15 @@ pub struct Coordinator<B> {
     token: CancellationToken,
 }
 
+/// Runs things.
+///
+/// # Liveness concerns
+///
+/// If you use one of the streaming versions (e.g. `begin_execute`),
+/// you need to make sure that the stdout / stderr / status channels
+/// are continuously read from or dropped completely. If not, one
+/// channel can fill up, preventing the other channels from receiving
+/// data as well.
 impl<B> Coordinator<B>
 where
     B: Backend,
@@ -3188,33 +3213,35 @@ mod tests {
         let ActiveExecution {
             task,
             stdin_tx: _,
-            mut stdout_rx,
+            stdout_rx,
             stderr_rx,
             status_rx: _,
         } = coordinator
             .begin_execute(token.clone(), request)
             .await
             .unwrap();
 
-        // Wait for some output before killing
-        let early_stdout = stdout_rx.recv().with_timeout().await.unwrap();
+        let stdout_rx = ReceiverStream::new(stdout_rx);
+        let stderr_rx = ReceiverStream::new(stderr_rx);
 
-        token.cancel();
+        // We (a) want to wait for some output before we try to
+        // kill the process and (b) need to keep pumping stdout /
+        // stderr / status to avoid locking up the output.
+        let stdout_rx = stdout_rx.inspect(|_| token.cancel());
 
         let WithOutput {
             response,
             stdout,
             stderr,
-        } = WithOutput::try_absorb(task, stdout_rx, stderr_rx)
+        } = WithOutput::try_absorb_stream(task, stdout_rx, stderr_rx)
             .with_timeout()
             .await
             .unwrap();
 
         assert!(!response.success, "{stderr}");
         assert_contains!(response.exit_detail, "kill");
 
-        assert_contains!(early_stdout, "Before");
-        assert_not_contains!(stdout, "Before");
+        assert_contains!(stdout, "Before");
         assert_not_contains!(stdout, "After");
 
         coordinator.shutdown().await?;