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

rust/private/rust.bzl

@@ -953,8 +953,8 @@ def _common_attrs_for_binary_without_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,
@@ -970,6 +970,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(_common_attrs_for_binary_without_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,

rust/repositories.bzl

@@ -59,6 +59,16 @@ def rules_rust_dependencies():
         url = "https://github.com/bazelbuild/apple_support/releases/download/0.11.0/apple_support.0.11.0.tar.gz",
     )
 
+    # process_wrapper needs a low-dependency way to process json.
+    maybe(
+        http_archive,
+        name = "rules_rust_tinyjson",
+        sha256 = "9c21866c7f051ebcefd028996494a374b7408ef946826cefc9761d58cce0fd36",
+        url = "https://github.com/rhysd/tinyjson/archive/refs/tags/v2.3.0.zip",
+        strip_prefix = "tinyjson-2.3.0",
+        build_file = "@rules_rust//util/process_wrapper:BUILD.tinyjson.bazel",
+    )
+
 # buildifier: disable=unnamed-macro
 def rust_register_toolchains(
         dev_components = False,

test/process_wrapper/BUILD.bazel

@@ -1,6 +1,7 @@
 load("@bazel_skylib//rules:build_test.bzl", "build_test")
 load("@bazel_skylib//rules:diff_test.bzl", "diff_test")
 load("@rules_cc//cc:defs.bzl", "cc_binary")
+load("//rust:defs.bzl", "rust_binary", "rust_test")
 load("//test/process_wrapper:process_wrapper_tester.bzl", "process_wrapper_tester")
 
 cc_binary(
@@ -148,3 +149,18 @@ build_test(
         ":process_wrapper_combined",
     ],
 )
+
+rust_binary(
+    name = "fake_rustc",
+    srcs = ["fake_rustc.rs"],
+)
+
+rust_test(
+    name = "rustc_quit_on_rmeta",
+    srcs = ["rustc_quit_on_rmeta.rs"],
+    data = [
+        ":fake_rustc",
+        "//util/process_wrapper",
+    ],
+    deps = ["//tools/runfiles"],
+)

test/process_wrapper/fake_rustc.rs

@@ -0,0 +1,8 @@
+//! This binary mocks the output of rustc when run with `--error-format=json` and `--json=artifacts`.
+
+fn main() {
+    eprintln!(r#"{{"rendered": "should be\nin output"}}"#);
+    eprintln!(r#"{{"emit": "metadata"}}"#);
+    std::thread::sleep(std::time::Duration::from_secs(1));
+    eprintln!(r#"{{"rendered": "should not be in output"}}"#);
+}

test/process_wrapper/rustc_quit_on_rmeta.rs

@@ -0,0 +1,94 @@
+#[cfg(test)]
+mod test {
+    use std::path::PathBuf;
+    use std::process::Command;
+    use std::str;
+
+    use runfiles::Runfiles;
+
+    // fake_rustc runs the fake_rustc binary under process_wrapper with the specified
+    // process wrapper arguments. No arguments are passed to fake_rustc itself.
+    fn fake_rustc(process_wrapper_args: &[&'static str]) -> String {
+        let r = Runfiles::create().unwrap();
+        let fake_rustc = r.rlocation(
+            &[
+                "rules_rust",
+                "test",
+                "process_wrapper",
+                if cfg!(unix) {
+                    "fake_rustc"
+                } else {
+                    "fake_rustc.exe"
+                },
+            ]
+            .iter()
+            .collect::<PathBuf>(),
+        );
+
+        let process_wrapper = r.rlocation(
+            &[
+                "rules_rust",
+                "util",
+                "process_wrapper",
+                if cfg!(unix) {
+                    "process_wrapper"
+                } else {
+                    "process_wrapper.exe"
+                },
+            ]
+            .iter()
+            .collect::<PathBuf>(),
+        );
+
+        let output = Command::new(process_wrapper)
+            .args(process_wrapper_args)
+            .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(),
+        );
+
+        String::from_utf8(output.stderr).unwrap()
+    }
+
+    #[test]
+    fn test_rustc_quit_on_rmeta_quits() {
+        let out_content = fake_rustc(&["--rustc-quit-on-rmeta", "true"]);
+        assert!(
+            !out_content.contains("should not be in output"),
+            "output should not contain 'should not be in output' but did: {}",
+            out_content
+        );
+    }
+
+    #[test]
+    fn test_rustc_quit_on_rmeta_output_json() {
+        let json_content = fake_rustc(&[
+            "--rustc-quit-on-rmeta",
+            "true",
+            "--rustc-output-format",
+            "json",
+        ]);
+        assert_eq!(
+            json_content,
+            concat!(r#"{"rendered": "should be\nin output"}"#, "\n")
+        );
+    }
+
+    #[test]
+    fn test_rustc_quit_on_rmeta_output_rendered() {
+        let rendered_content = fake_rustc(&[
+            "--rustc-quit-on-rmeta",
+            "true",
+            "--rustc-output-format",
+            "rendered",
+        ]);
+        assert_eq!(rendered_content, "should be\nin output");
+    }
+}

util/process_wrapper/BUILD.bazel

@@ -8,6 +8,9 @@ rust_binary_without_process_wrapper(
     srcs = glob(["*.rs"]),
     edition = "2018",
     visibility = ["//visibility:public"],
+    deps = [
+        "@rules_rust_tinyjson//:tinyjson",
+    ],
 )
 
 rust_test(

util/process_wrapper/BUILD.tinyjson.bazel

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

util/process_wrapper/main.rs

@@ -14,50 +14,108 @@
 
 mod flags;
 mod options;
+mod output;
+mod rustc;
 mod util;
 
 use std::fs::{copy, OpenOptions};
-use std::process::{exit, Command, Stdio};
+use std::io;
+use std::process::{exit, Command, ExitStatus, Stdio};
 
 use crate::options::options;
+use crate::output::{process_output, LineOutput};
+
+#[cfg(windows)]
+fn status_code(status: ExitStatus, was_killed: bool) -> i32 {
+    // On windows, there's no good way to know if the process was killed by a signal.
+    // If we killed the process, we override the code to signal success.
+    if was_killed {
+        0
+    } else {
+        status.code().unwrap_or(1)
+    }
+}
+
+#[cfg(not(windows))]
+fn status_code(status: ExitStatus, was_killed: bool) -> i32 {
+    // On unix, if code is None it means that the process was killed by a signal.
+    // https://doc.rust-lang.org/std/process/struct.ExitStatus.html#method.success
+    match status.code() {
+        Some(code) => code,
+        // If we killed the process, we expect None here
+        None if was_killed => 0,
+        // Otherwise it's some unexpected signal
+        None => 1,
+    }
+}
 
 fn main() {
     let opts = match options() {
         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()
+
+    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 {
-        Stdio::inherit()
+        Box::new(io::stderr())
     };
-    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(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()
+        })
+        .stderr(Stdio::piped())
+        .spawn()
         .expect("process wrapper error: failed to spawn child process");
 
-    if status.success() {
+    let child_stderr = Box::new(child.stderr.take().unwrap());
+
+    let mut was_killed = false;
+    let result = if !opts.rustc_quit_on_rmeta {
+        // Process output normally by forwarding stderr
+        process_output(child_stderr, stderr, LineOutput::Message)
+    } else {
+        let format = opts.rustc_output_format;
+        let mut kill = false;
+        let result = process_output(child_stderr, stderr, |line| {
+            rustc::stop_on_rmeta_completion(line, format, &mut kill)
+        });
+        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;
+        }
+        result
+    };
+    result.expect("process wrapper error: failed to process stderr");
+
+    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 = status_code(status, was_killed);
+    let success = code == 0;
+    if success {
         if let Some(tf) = opts.touch_file {
             OpenOptions::new()
                 .create(true)
@@ -75,5 +133,5 @@ fn main() {
         }
     }
 
-    exit(status.code().unwrap())
+    exit(code)
 }