Prevent massive amounts of outgoing stderr / stdout

We've noticed some bad resource usage patterns in production which
manifests as 100% CPU usage, split evenly between system and user.

My hypothesis is that people have written some infinite loop that
prints output. We are able to transfer that data from the process to the
worker to the coordinator/ui to nginx to the user. However, the user
can't actually consume 10+ MiB of output.

There's also a possibility that Docker has effectively an infinite
buffer for stdout/stderr. In that case, we might just be bloating up
memory somewhere.

This is intended as a first step, to find people that are legitimately
trying to produce a bunch of output. From there, we could potentially
tweak the code to allow people to opt-in to generating even more data.

Author: shepmaster

compiler/base/orchestrator/src/coordinator.rs

@@ -4055,6 +4055,39 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    #[snafu::report]
+    async fn amount_of_output_is_limited() -> Result<()> {
+        // The limits are only applied to the container
+        let coordinator = new_coordinator_docker().await;
+
+        let req = ExecuteRequest {
+            code: r##"
+                use std::io::Write;
+
+                fn main() {
+                    let a = "a".repeat(1024);
+                    let out = std::io::stdout();
+                    let mut out = out.lock();
+                    loop {//for _ in 0..1_000_000 {
+                        let _ = out.write_all(a.as_bytes());
+                        let _ = out.write_all(b"\n");
+                    }
+                }
+            "##
+            .into(),
+            ..new_execution_limited_request()
+        };
+
+        let err = coordinator.execute(req).with_timeout().await.unwrap_err();
+        let err = snafu::ChainCompat::new(&err).last().unwrap();
+        assert_contains!(err.to_string(), "bytes of output, exiting");
+
+        coordinator.shutdown().await?;
+
+        Ok(())
+    }
+
     static TIMEOUT: Lazy<Duration> = Lazy::new(|| {
         let millis = env::var("TESTS_TIMEOUT_MS")
             .ok()

compiler/base/orchestrator/src/worker.rs

@@ -612,24 +612,43 @@ async fn process_end(
 ) -> Result<ExecuteCommandResponse, ProcessError> {
     use process_error::*;
 
-    select! {
-        () = token.cancelled() => child.kill().await.context(KillChildSnafu)?,
-        _ = child.wait() => {},
+    let mut killed = false;
+
+    let status = loop {
+        select! {
+            // The user requested that the process be killed
+            () = token.cancelled(), if !killed => {
+                child.kill().await.context(KillChildSnafu)?;
+                killed = true;
+            },
+
+            // The process exited normally
+            status = child.wait() => break status,
+
+            // One of our tasks exited unexpectedly
+            // TODO: dedupe errors or fully split them
+            Some(task) = task_set.join_next() => {
+                task.context(StdioTaskPanickedSnafu)?
+                    .context(StdioTaskFailedSnafu)?;
+            },
+        };
     };
 
-    let status = child.wait().await.context(WaitChildSnafu)?;
+    let status = status.context(WaitChildSnafu)?;
 
     stdin_shutdown_tx
         .send(job_id)
         .await
         .drop_error_details()
         .context(UnableToSendStdinShutdownSnafu)?;
 
+    // Check any remaining tasks to see if they had an error
     while let Some(task) = task_set.join_next().await {
         task.context(StdioTaskPanickedSnafu)?
             .context(StdioTaskFailedSnafu)?;
     }
 
+    // TODO: check this for death earlier?
     statistics_task
         .await
         .context(StatisticsTaskPanickedSnafu)?
@@ -1213,6 +1232,8 @@ mod test {
     }
 }
 
+const OUTPUT_BYTE_LIMIT: usize = 640 * 1024;
+
 async fn copy_child_output(
     output: impl AsyncRead + Unpin,
     coordinator_tx: MultiplexingSender,
@@ -1221,17 +1242,28 @@ async fn copy_child_output(
     use copy_child_output_error::*;
 
     let mut buf = Utf8BufReader::new(output);
+    let mut n_total_bytes: usize = 0;
 
     while let Some(buffer) = buf.next().await.context(UnableToReadSnafu)? {
+        let n_bytes = buffer.len();
+
         coordinator_tx
             .send_ok(xform(buffer))
             .await
             .context(UnableToSendSnafu)?;
+
+        n_total_bytes = n_total_bytes.saturating_add(n_bytes);
+        ensure!(
+            n_total_bytes <= OUTPUT_BYTE_LIMIT,
+            TooManyBytesSnafu { n_total_bytes }
+        );
     }
 
     Ok(())
 }
 
+const BYTE_LIMIT_URL: &str = "https://github.com/rust-lang/rust-playground/discussions/1027";
+
 #[derive(Debug, Snafu)]
 #[snafu(module)]
 pub enum CopyChildOutputError {
@@ -1240,6 +1272,11 @@ pub enum CopyChildOutputError {
 
     #[snafu(display("Failed to send output packet"))]
     UnableToSend { source: MultiplexingSenderError },
+
+    #[snafu(display(
+        "Generated {n_total_bytes} bytes of output, exiting (640K ought to be enough for anybody). If this was not an accident, tell us more at {BYTE_LIMIT_URL}"
+    ))]
+    TooManyBytes { n_total_bytes: usize },
 }
 
 // stdin/out <--> messages.