process_wrapper: make process_wrapper not use extra threads

Author: gigaroby

util/process_wrapper/main.rs

@@ -21,11 +21,9 @@ mod util;
 use std::fs::{copy, OpenOptions};
 use std::io;
 use std::process::{exit, Command, ExitStatus, Stdio};
-use std::sync::mpsc::sync_channel;
-
-use output::process_output;
 
 use crate::options::options;
+use crate::output::{process_output, LineOutput};
 
 #[cfg(windows)]
 fn status_code(status: ExitStatus, was_killed: bool) -> i32 {
@@ -92,27 +90,24 @@ fn main() {
     let child_stderr = Box::new(child.stderr.take().unwrap());
 
     let mut was_killed = false;
-    let stderr_thread = if !opts.rustc_quit_on_rmeta {
+    let result = if !opts.rustc_quit_on_rmeta {
         // Process output normally by forwarding stderr
-        process_output(child_stderr, stderr, Some)
+        process_output(child_stderr, stderr, LineOutput::Message)
     } else {
         let format = opts.rustc_output_format;
-        // Process json rustc output and kill the subprocess when we get a signal
-        // that we emitted a metadata file.
-        // This receiver will block until a corresponding send happens.
-        let (stop_sender, stop) = sync_channel(0);
-        let thread = process_output(child_stderr, stderr, move |line| {
-            rustc::stop_on_rmeta_completion(line, format, &stop_sender)
+        let mut kill = false;
+        let result = process_output(child_stderr, stderr, |line| {
+            rustc::stop_on_rmeta_completion(line, format, &mut kill)
         });
-        if stop.recv().is_ok() {
+        if kill {
             // If recv returns Ok(), a signal was sent in this channel so we should terminate the child process.
             // We can safely ignore the Result from kill() as we don't care if the process already terminated.
             let _ = child.kill();
             was_killed = true;
         }
-        thread
+        result
     };
-    stderr_thread.join().unwrap().unwrap();
+    result.expect("process wrapper error: failed to process stderr");
 
     let status = child
         .wait()

util/process_wrapper/output.rs

@@ -1,27 +1,56 @@
+// Copyright 2020 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 use std::io::{self, prelude::*};
-use std::thread;
 
+/// LineOutput tells process_output what to do when a line is processed.
+/// If a Message is returned, it will be written to write_end, if
+/// Skip is returned nothing will be printed and execution continues,
+/// if Terminate is returned, process_output returns immediately.
+/// Terminate is used to stop processing when we see an emit metadata
+/// message.
+#[derive(Debug)]
+pub(crate) enum LineOutput {
+    Message(String),
+    Skip,
+    Terminate,
+}
+
+/// process_output reads lines from read_end and invokes process_line on each.
+/// Depending on the result of process_line, the modified message may be written
+/// to write_end.
 pub(crate) fn process_output<F>(
-    read_end: Box<dyn Read + Send>,
-    write_end: Box<dyn Write + Send>,
+    read_end: Box<dyn Read>,
+    write_end: Box<dyn Write>,
     mut process_line: F,
-) -> thread::JoinHandle<io::Result<()>>
+) -> io::Result<()>
 where
-    F: FnMut(String) -> Option<String> + Send + 'static,
+    F: FnMut(String) -> LineOutput,
 {
-    thread::spawn(move || {
-        let mut reader = io::BufReader::new(read_end);
-        let mut writer = io::LineWriter::new(write_end);
-        loop {
-            let mut line = String::new();
-            let read_bytes = reader.read_line(&mut line)?;
-            if read_bytes == 0 {
-                break;
-            }
-            if let Some(to_write) = process_line(line) {
-                writer.write_all(to_write.as_bytes())?;
-            }
+    let mut reader = io::BufReader::new(read_end);
+    let mut writer = io::LineWriter::new(write_end);
+    loop {
+        let mut line = String::new();
+        let read_bytes = reader.read_line(&mut line)?;
+        if read_bytes == 0 {
+            break;
         }
-        Ok(())
-    })
+        match process_line(line) {
+            LineOutput::Message(to_write) => writer.write_all(to_write.as_bytes())?,
+            LineOutput::Skip => {}
+            LineOutput::Terminate => return Ok(()),
+        };
+    }
+    Ok(())
 }

util/process_wrapper/rustc.rs

@@ -1,6 +1,21 @@
-use std::sync::mpsc::SyncSender;
+// Copyright 2020 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 use tinyjson::JsonValue;
 
+use crate::output::LineOutput;
+
 #[derive(Debug, Copy, Clone)]
 pub(crate) enum ErrorFormat {
     Json,
@@ -25,36 +40,39 @@ fn get_key(value: &JsonValue, key: &str) -> Option<String> {
     }
 }
 
-// stop_on_rmeta_completion takes an output line from rustc configured with
-// --error-format=json, parses the json and returns the appropriate output
-// according to the original --error-format supplied to rustc.
-// In addition, it will send a signal to the stop channel when metadata is emitted
-// so the compiler can be terminated.
-// This is used to implement pipelining in rules_rust, please see
-// https://internals.rust-lang.org/t/evaluating-pipelined-rustc-compilation/10199
+/// stop_on_rmeta_completion takes an output line from rustc configured with
+/// --error-format=json, parses the json and returns the appropriate output
+/// according to the original --error-format supplied to rustc.
+/// In addition, it will signal to stop when metadata is emitted
+/// so the compiler can be terminated.
+/// This is used to implement pipelining in rules_rust, please see
+/// https://internals.rust-lang.org/t/evaluating-pipelined-rustc-compilation/10199
 pub(crate) fn stop_on_rmeta_completion(
     line: String,
     error_format: ErrorFormat,
-    stop: &SyncSender<()>,
-) -> Option<String> {
+    kill: &mut bool,
+) -> LineOutput {
     let parsed: JsonValue = line
         .parse()
         .expect("process wrapper error: expected json messages in pipeline mode");
     if let Some(emit) = get_key(&parsed, "emit") {
         // We don't want to print emit messages.
         // If the emit messages is "metadata" we can signal the process to quit
-        if emit == "metadata" {
-            stop.send(())
-                .expect("process wrapper error: receiver closed");
-        }
-        return None;
+        return if emit == "metadata" {
+            *kill = true;
+            LineOutput::Terminate
+        } else {
+            LineOutput::Skip
+        };
     };
 
     match error_format {
         // If the output should be json, we just forward the messages as-is
-        ErrorFormat::Json => Some(line),
+        ErrorFormat::Json => LineOutput::Message(line),
         // Otherwise we extract the "rendered" attribute.
         // If we don't find it we skip the line.
-        _ => get_key(&parsed, "rendered"),
+        _ => get_key(&parsed, "rendered")
+            .map(LineOutput::Message)
+            .unwrap_or(LineOutput::Skip),
     }
 }