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

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)
 }