process_wrapper: add support for terminating rustc after it emits rmeta.

This is a key component of supporting pipelining in rules_rust. The
bazel rules will (in a follow-up PR) configure rustc to run either to
completion for rules whose dependencies require the full rlib files or
until they emit the rmeta files if dependencies only require that.

This is safe to commit as there are no changes to user-visible behavior
until the new flags are used.

Author: gigaroby

WORKSPACE.bazel

@@ -53,6 +53,14 @@ http_archive(
     url = "https://github.com/bazelbuild/continuous-integration/releases/download/rules-1.0.0/bazelci_rules-1.0.0.tar.gz",
 )
 
+http_archive(
+    name = "tinyjson",
+    build_file = "@//util/process_wrapper:BUILD.tinyjson.bazel",
+    sha256 = "9c21866c7f051ebcefd028996494a374b7408ef946826cefc9761d58cce0fd36",
+    strip_prefix = "tinyjson-2.3.0",
+    url = "https://github.com/rhysd/tinyjson/archive/refs/tags/v2.3.0.zip",
+)
+
 # To run with RBE on Bazel CI, uncomment the following lines.
 #
 # load("@bazelci_rules//:rbe_repo.bzl", "rbe_preconfig")

rust/private/rust.bzl

@@ -939,8 +939,8 @@ def _fake_out_process_wrapper(attrs):
     )
     return new_attr
 
-# Provides an internal rust_binary to use that we can use to build the process
-# wrapper, this breaks the dependency of rust_binary on the process wrapper by
+# Provides an internal rust_{binary,library} to use that we can use to build the process
+# wrapper, this breaks the dependency of rust_* on the process wrapper by
 # setting it to None, which the functions in rustc detect and build accordingly.
 rust_binary_without_process_wrapper = rule(
     implementation = _rust_binary_impl,
@@ -956,6 +956,19 @@ rust_binary_without_process_wrapper = rule(
     incompatible_use_toolchain_transition = True,
 )
 
+rust_library_without_process_wrapper = rule(
+    implementation = _rust_library_impl,
+    provides = _common_providers,
+    attrs = dict(_fake_out_process_wrapper(_common_attrs).items()),
+    fragments = ["cpp"],
+    host_fragments = ["cpp"],
+    toolchains = [
+        str(Label("//rust:toolchain")),
+        "@bazel_tools//tools/cpp:toolchain_type",
+    ],
+    incompatible_use_toolchain_transition = True,
+)
+
 rust_test = rule(
     implementation = _rust_test_impl,
     provides = _common_providers,

test/rustc_quit_on_rmeta/BUILD.bazel

@@ -0,0 +1,16 @@
+load("//rust:defs.bzl", "rust_test")
+
+sh_binary(
+    name = "fake_rustc",
+    srcs = ["fake_rustc.sh"],
+)
+
+rust_test(
+    name = "rustc_quit_on_rmeta",
+    srcs = ["rustc_quit_on_rmeta.rs"],
+    data = [
+        ":fake_rustc",
+        "//util/process_wrapper",
+    ],
+    deps = ["//tools/runfiles"],
+)

test/rustc_quit_on_rmeta/fake_rustc.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+# Redirect stdout to stderr
+exec >&2
+
+echo '{"rendered": "I am a fake rustc\nvery very fake"}'
+echo '{"emit": "metadata"}'
+sleep 1s
+echo '{"rendered": "I should not print this"}'

test/rustc_quit_on_rmeta/rustc_quit_on_rmeta.rs

@@ -0,0 +1,45 @@
+use std::env;
+use std::fs::read_to_string;
+use std::path::Path;
+use std::process::Command;
+use std::str;
+
+use runfiles::Runfiles;
+
+fn fake_rustc(process_wrapper_args: &[&'static str]) ->  String {
+    let r = Runfiles::create().unwrap();
+    let fake_rustc = r.rlocation("rules_rust/test/rustc_quit_on_rmeta/fake_rustc");
+    let process_wrapper = r.rlocation("rules_rust/util/process_wrapper/process_wrapper");
+    let stderr_file_path = Path::new(&env::var("TEST_TMPDIR").unwrap()).join("stdout.txt");
+
+    let output = Command::new(process_wrapper)
+        .args(process_wrapper_args)
+        .arg("--stderr-file")
+        .arg(&stderr_file_path)
+        .arg("--")
+        .arg(fake_rustc)
+        .output().unwrap();
+
+    assert!(
+        output.status.success(),
+        "unable to run process_wrapper: {} {}",
+        str::from_utf8(&output.stdout).unwrap(),
+        str::from_utf8(&output.stderr).unwrap(),
+    );
+
+    read_to_string(stderr_file_path).unwrap()
+}
+
+#[test]
+fn test_rustc_quit_on_rmeta_quits() {
+    let out_content = fake_rustc(&["--rustc-quit-on-rmeta", "true"]);
+    assert!(!out_content.contains("I should not print this"), "output should not contain 'I should not print this' but did: {}", out_content);
+}
+
+#[test]
+fn test_rustc_quit_on_rmeta_output_format() {
+    let json_content = fake_rustc(&["--rustc-quit-on-rmeta", "true", "--rustc-output-format", "json"]);
+    assert_eq!(json_content, concat!(r#"{"rendered": "I am a fake rustc\nvery very fake"}"#, "\n"));
+    let rendered_content = fake_rustc(&["--rustc-quit-on-rmeta", "true", "--rustc-output-format", "rendered"]);
+    assert_eq!(rendered_content, "I am a fake rustc\nvery very fake\n");
+}

util/process_wrapper/BUILD.bazel

@@ -7,6 +7,9 @@ rust_binary_without_process_wrapper(
     name = "process_wrapper",
     srcs = glob(["*.rs"]),
     visibility = ["//visibility:public"],
+    deps = [
+        "@tinyjson",
+    ],
 )
 
 rust_test(

util/process_wrapper/BUILD.tinyjson.bazel

@@ -0,0 +1,8 @@
+# buildifier: disable=bzl-visibility
+load("@//rust/private:rust.bzl", "rust_library_without_process_wrapper")
+
+rust_library_without_process_wrapper(
+    name = "tinyjson",
+    srcs = glob(["src/*.rs"]),
+    visibility = ["@//util/process_wrapper:__pkg__"],
+)

util/process_wrapper/main.rs

@@ -14,10 +14,16 @@
 
 mod flags;
 mod options;
+mod output;
+mod rustc;
 mod util;
 
 use std::fs::{copy, OpenOptions};
+use std::io;
 use std::process::{exit, Command, Stdio};
+use std::sync::mpsc::sync_channel;
+
+use output::process_output;
 
 use crate::options::options;
 
@@ -26,38 +32,83 @@ fn main() {
         Err(err) => panic!("process wrapper error: {}", err),
         Ok(v) => v,
     };
-    let stdout = if let Some(stdout_file) = opts.stdout_file {
-        OpenOptions::new()
-            .create(true)
-            .truncate(true)
-            .write(true)
-            .open(stdout_file)
-            .expect("process wrapper error: unable to open stdout file")
-            .into()
-    } else {
-        Stdio::inherit()
-    };
-    let stderr = if let Some(stderr_file) = opts.stderr_file {
-        OpenOptions::new()
-            .create(true)
-            .truncate(true)
-            .write(true)
-            .open(stderr_file)
-            .expect("process wrapper error: unable to open stderr file")
-            .into()
-    } else {
-        Stdio::inherit()
-    };
-    let status = Command::new(opts.executable)
+
+    let mut child = Command::new(opts.executable)
         .args(opts.child_arguments)
         .env_clear()
         .envs(opts.child_environment)
-        .stdout(stdout)
-        .stderr(stderr)
-        .status()
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()
         .expect("process wrapper error: failed to spawn child process");
 
-    if status.success() {
+    let stdout: Box<dyn io::Write + Send> = if let Some(stdout_file) = opts.stdout_file {
+        Box::new(
+            OpenOptions::new()
+                .create(true)
+                .truncate(true)
+                .write(true)
+                .open(stdout_file)
+                .expect("process wrapper error: unable to open stdout file"),
+        )
+    } else {
+        Box::new(io::stdout())
+    };
+    let stderr: Box<dyn io::Write + Send> = if let Some(stderr_file) = opts.stderr_file {
+        Box::new(
+            OpenOptions::new()
+                .create(true)
+                .truncate(true)
+                .write(true)
+                .open(stderr_file)
+                .expect("process wrapper error: unable to open stderr file"),
+        )
+    } else {
+        Box::new(io::stderr())
+    };
+
+    let child_stdout = Box::new(child.stdout.take().unwrap());
+    let child_stderr = Box::new(child.stderr.take().unwrap());
+
+    if !opts.rustc_quit_on_rmeta {
+        // Process output normally by forwarding stdout and stderr
+        let stdout_thread = process_output(child_stdout, stdout, |line| Some(line));
+        let stderr_thread = process_output(child_stderr, stderr, |line| Some(line));
+        stdout_thread.join().unwrap().unwrap();
+        stderr_thread.join().unwrap().unwrap();
+    } 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_stdout, stop) = sync_channel(0);
+        let stop_sender_stderr = stop_sender_stdout.clone();
+        let stdout_thread = process_output(child_stdout, stdout, move |line| {
+            rustc::process_message(line, format, &stop_sender_stdout)
+        });
+        let stderr_thread = process_output(child_stderr, stderr, move |line| {
+            rustc::process_message(line, format, &stop_sender_stderr)
+        });
+        if let Ok(_) = stop.recv() {
+            // 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();
+        }
+
+        stdout_thread.join().unwrap().unwrap();
+        stderr_thread.join().unwrap().unwrap();
+    }
+
+    let status = child
+        .wait()
+        .expect("process wrapper error: failed to wait for child process");
+    // If the child process is rustc and is killed after metadata generation, that's also a success.
+    let (code, success) = match status.code() {
+        // According to rust docs: Signal termination is not considered a success, and success is defined as a zero exit status.
+        Some(code) => (code, code == 0),
+        None => (0, true),
+    };
+    if success {
         if let Some(tf) = opts.touch_file {
             OpenOptions::new()
                 .create(true)
@@ -75,5 +126,5 @@ fn main() {
         }
     }
 
-    exit(status.code().unwrap())
+    exit(code)
 }

util/process_wrapper/options.rs

@@ -4,6 +4,7 @@ use std::fmt;
 use std::process::exit;
 
 use crate::flags::{FlagParseError, Flags, ParseOutcome};
+use crate::rustc;
 use crate::util::*;
 
 #[derive(Debug)]
@@ -38,6 +39,12 @@ pub(crate) struct Options {
     pub(crate) stdout_file: Option<String>,
     // If set, redirects the child process stderr to this file.
     pub(crate) stderr_file: Option<String>,
+    // If set, it configures rustc to emit an rmeta file and then
+    // quit.
+    pub(crate) rustc_quit_on_rmeta: bool,
+    // If rustc_quit_on_rmeta is set to true, this controls the
+    // output format of rustc messages.
+    pub(crate) rustc_output_format: rustc::Output,
 }
 
 pub(crate) fn options() -> Result<Options, OptionError> {
@@ -51,6 +58,8 @@ pub(crate) fn options() -> Result<Options, OptionError> {
     let mut copy_output_raw = None;
     let mut stdout_file = None;
     let mut stderr_file = None;
+    let mut rustc_quit_on_rmeta_raw = None;
+    let mut rustc_output_format_raw = None;
     let mut flags = Flags::new();
     flags.define_repeated_flag("--subst", "", &mut subst_mapping_raw);
     flags.define_flag("--volatile-status-file", "", &mut volatile_status_file_raw);
@@ -80,6 +89,19 @@ pub(crate) fn options() -> Result<Options, OptionError> {
         "Redirect subprocess stderr in this file.",
         &mut stderr_file,
     );
+    flags.define_flag(
+        "--rustc-quit-on-rmeta",
+        "If enabled, this wrapper will terminate rustc after rmeta has been emitted.",
+        &mut rustc_quit_on_rmeta_raw,
+    );
+    flags.define_flag(
+        "--rustc-output-format",
+        "Controls the rustc output format if --rustc-quit-on-rmeta is set.\n\
+        'json' will cause the json output to be output, \
+        'rendered' will extract the rendered message and print that.\n\
+        Default: `rendered`",
+        &mut rustc_output_format_raw,
+    );
 
     let mut child_args = match flags
         .parse(env::args().collect())
@@ -138,6 +160,19 @@ pub(crate) fn options() -> Result<Options, OptionError> {
         })
         .transpose()?;
 
+    let rustc_quit_on_rmeta = rustc_quit_on_rmeta_raw.map_or(false, |s| s == "true");
+    let rustc_output_format = rustc_output_format_raw
+        .map(|v| match v.as_str() {
+            "json" => Ok(rustc::Output::Json),
+            "rendered" => Ok(rustc::Output::Rendered),
+            _ => Err(OptionError::Generic(format!(
+                "invalid --rustc-output-format '{}'",
+                v
+            ))),
+        })
+        .transpose()?
+        .unwrap_or_default();
+
     // Prepare the environment variables, unifying those read from files with the ones
     // of the current process.
     let vars = environment_block(environment_file_block, &stamp_mappings, &subst_mappings);
@@ -159,6 +194,8 @@ pub(crate) fn options() -> Result<Options, OptionError> {
         copy_output,
         stdout_file,
         stderr_file,
+        rustc_quit_on_rmeta,
+        rustc_output_format,
     })
 }
 

util/process_wrapper/output.rs

@@ -0,0 +1,27 @@
+use std::io::{self, prelude::*};
+use std::thread;
+
+pub(crate) fn process_output<F>(
+    read_end: Box<dyn Read + Send>,
+    write_end: Box<dyn Write + Send>,
+    mut process_line: F,
+) -> thread::JoinHandle<io::Result<()>>
+where
+    F: FnMut(String) -> Option<String> + Send + 'static,
+{
+    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.into_bytes())?;
+            }
+        }
+        Ok(())
+    })
+}