diff --git a/rust/private/rust.bzl b/rust/private/rust.bzl index 6a2ed627f1..963dfea026 100644 --- a/rust/private/rust.bzl +++ b/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, diff --git a/rust/repositories.bzl b/rust/repositories.bzl index 99f2f51ed3..957dd2c226 100644 --- a/rust/repositories.bzl +++ b/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 = "@rules_rust//util/process_wrapper:BUILD.tinyjson.bazel", + ) + # buildifier: disable=unnamed-macro def rust_register_toolchains( dev_components = False, diff --git a/test/process_wrapper/BUILD.bazel b/test/process_wrapper/BUILD.bazel index a9f2189b87..0a3e1bcffe 100644 --- a/test/process_wrapper/BUILD.bazel +++ b/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"], +) diff --git a/test/process_wrapper/fake_rustc.rs b/test/process_wrapper/fake_rustc.rs new file mode 100755 index 0000000000..a82a82bc9c --- /dev/null +++ b/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"}}"#); +} diff --git a/test/process_wrapper/rustc_quit_on_rmeta.rs b/test/process_wrapper/rustc_quit_on_rmeta.rs new file mode 100644 index 0000000000..0e0b42f32a --- /dev/null +++ b/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::(), + ); + + let process_wrapper = r.rlocation( + &[ + "rules_rust", + "util", + "process_wrapper", + if cfg!(unix) { + "process_wrapper" + } else { + "process_wrapper.exe" + }, + ] + .iter() + .collect::(), + ); + + 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"); +} diff --git a/util/process_wrapper/BUILD.bazel b/util/process_wrapper/BUILD.bazel index ebceb4a5fd..8f32cb2f6c 100644 --- a/util/process_wrapper/BUILD.bazel +++ b/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( diff --git a/util/process_wrapper/BUILD.tinyjson.bazel b/util/process_wrapper/BUILD.tinyjson.bazel new file mode 100644 index 0000000000..6e2f716bf6 --- /dev/null +++ b/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__"], +) diff --git a/util/process_wrapper/main.rs b/util/process_wrapper/main.rs index 41140a37e2..fe4cc33060 100644 --- a/util/process_wrapper/main.rs +++ b/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 = 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 = 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) } diff --git a/util/process_wrapper/options.rs b/util/process_wrapper/options.rs index 24bba9fe80..7f1ab124be 100644 --- a/util/process_wrapper/options.rs +++ b/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, // If set, redirects the child process stderr to this file. pub(crate) stderr_file: Option, + // 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 { @@ -51,6 +58,8 @@ pub(crate) fn options() -> Result { 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 { "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 { }) .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 { copy_output, stdout_file, stderr_file, + rustc_quit_on_rmeta, + rustc_output_format, }) } diff --git a/util/process_wrapper/output.rs b/util/process_wrapper/output.rs new file mode 100644 index 0000000000..9040881051 --- /dev/null +++ b/util/process_wrapper/output.rs @@ -0,0 +1,27 @@ +use std::io::{self, prelude::*}; +use std::thread; + +pub(crate) fn process_output( + read_end: Box, + write_end: Box, + mut process_line: F, +) -> thread::JoinHandle> +where + F: FnMut(String) -> Option + 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(()) + }) +} diff --git a/util/process_wrapper/rustc.rs b/util/process_wrapper/rustc.rs new file mode 100644 index 0000000000..42422e217b --- /dev/null +++ b/util/process_wrapper/rustc.rs @@ -0,0 +1,53 @@ +use std::sync::mpsc::SyncSender; +use tinyjson::JsonValue; + +#[derive(Debug, Copy, Clone)] +pub(crate) enum Output { + Json, + Rendered, +} + +impl Default for Output { + fn default() -> Self { + Self::Rendered + } +} + +fn get_key(value: JsonValue, key: &str) -> Option { + if let JsonValue::Object(mut map) = value { + if let JsonValue::String(s) = map.remove(key)? { + Some(s) + } else { + None + } + } else { + None + } +} + +pub(crate) fn process_message( + line: String, + output: Output, + stop: &SyncSender<()>, +) -> Option { + let parsed: JsonValue = line + .parse() + .expect("process wrapper error: expected json messages in pipeline mode"); + if let Some(emit) = get_key(parsed.clone(), "emit") { + // We don't want to print emit messages. + // If the emit messages is "metadata" we can signal the process to quit + if emit == "metadata" { + stop.send(()) + .expect("process wrapper error: receiver closed"); + } + return None; + }; + + match output { + // If the output should be json, we just forward the messages as-is + Output::Json => Some(line), + // Otherwise we extract the "rendered" attribute. + // If we don't find it we skip the line. + _ => get_key(parsed, "rendered"), + } +}