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

@@ -70,3 +70,11 @@ http_archive(
 load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
 
 rules_pkg_dependencies()
+
+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",
+)

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,

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 = "tinyjson",
+        sha256 = "9c21866c7f051ebcefd028996494a374b7408ef946826cefc9761d58cce0fd36",
+        url = "https://github.com/rhysd/tinyjson/archive/refs/tags/v2.3.0.zip",
+        strip_prefix = "tinyjson-2.3.0",
+        build_file = "@//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,6 @@
+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,89 @@
+use std::path::PathBuf;
+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",
+            "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 = [
+        "@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,50 +14,126 @@
 
 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()
-    } 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());
+
+    let was_killed = if !opts.rustc_quit_on_rmeta {
+        // Process output normally by forwarding stdout and stderr
+        let stdout_thread = process_output(child_stdout, stdout, Some);
+        let stderr_thread = process_output(child_stderr, stderr, Some);
+        stdout_thread.join().unwrap().unwrap();
+        stderr_thread.join().unwrap().unwrap();
+        false
+    } else {
+        let mut was_killed = false;
+        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 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;
+        }
+
+        stdout_thread.join().unwrap().unwrap();
+        stderr_thread.join().unwrap().unwrap();
+        was_killed
+    };
+
+    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 +151,5 @@ fn main() {
         }
     }
 
-    exit(status.code().unwrap())
+    exit(code)
 }