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

@@ -952,8 +952,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,
@@ -969,6 +969,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,7 @@
+// This binary mocks the output of rustc when run with --error-format=json and --json=artifacts.
+fn main() {
+    eprintln!(r#"{{"rendered": "I am a fake rustc\nvery very fake"}}"#);
+    eprintln!(r#"{{"emit": "metadata"}}"#);
+    std::thread::sleep(std::time::Duration::from_secs(1));
+    eprintln!(r#"{{"rendered": "I should not print this"}}"#);
+}

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("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_json() {
+        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")
+        );
+    }
+
+    #[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, "I am a fake rustc\nvery very fake");
+    }
+}

util/process_wrapper/BUILD.bazel

@@ -7,6 +7,9 @@ rust_binary_without_process_wrapper(
     name = "process_wrapper",
     srcs = glob(["*.rs"]),
     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,113 @@
 
 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 std::sync::mpsc::sync_channel;
+
+use output::process_output;
 
 use crate::options::options;
 
+#[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 stderr_thread = if !opts.rustc_quit_on_rmeta {
+        // Process output normally by forwarding stderr
+        process_output(child_stderr, stderr, Some)
+    } 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)
+        });
+        if stop.recv().is_ok() {
+            // 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
+    };
+    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 = status_code(status, was_killed);
+    let success = code == 0;
+    if success {
         if let Some(tf) = opts.touch_file {
             OpenOptions::new()
                 .create(true)
@@ -75,5 +138,5 @@ fn main() {
         }
     }
 
-    exit(status.code().unwrap())
+    exit(code)
 }