From 189dfbc5508613fbd3571af0d961ed7a8475d819 Mon Sep 17 00:00:00 2001 From: Roberto Bampi Date: Wed, 23 Mar 2022 12:17:53 +0100 Subject: [PATCH] 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. --- rust/private/rust.bzl | 17 ++- rust/repositories.bzl | 10 ++ test/process_wrapper/BUILD.bazel | 16 +++ test/process_wrapper/fake_rustc.rs | 8 ++ test/process_wrapper/rustc_quit_on_rmeta.rs | 94 ++++++++++++++++ util/process_wrapper/BUILD.bazel | 3 + util/process_wrapper/BUILD.tinyjson.bazel | 8 ++ util/process_wrapper/main.rs | 112 +++++++++++++++----- util/process_wrapper/options.rs | 37 +++++++ util/process_wrapper/output.rs | 56 ++++++++++ util/process_wrapper/rustc.rs | 78 ++++++++++++++ 11 files changed, 410 insertions(+), 29 deletions(-) create mode 100755 test/process_wrapper/fake_rustc.rs create mode 100644 test/process_wrapper/rustc_quit_on_rmeta.rs create mode 100644 util/process_wrapper/BUILD.tinyjson.bazel create mode 100644 util/process_wrapper/output.rs create mode 100644 util/process_wrapper/rustc.rs diff --git a/rust/private/rust.bzl b/rust/private/rust.bzl index dfef5c88cf..48490eddbe 100644 --- a/rust/private/rust.bzl +++ b/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, diff --git a/rust/repositories.bzl b/rust/repositories.bzl index 99f2f51ed3..bff08604ab 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 = "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, 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..d7937c6cf2 --- /dev/null +++ b/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"}}"#); +} 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..df32341dc4 --- /dev/null +++ b/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::(), + ); + + 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("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"); + } +} diff --git a/util/process_wrapper/BUILD.bazel b/util/process_wrapper/BUILD.bazel index 3a2b2888b9..90d82490eb 100644 --- a/util/process_wrapper/BUILD.bazel +++ b/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( 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..a90bf4c3b8 100644 --- a/util/process_wrapper/main.rs +++ b/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 = 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) } diff --git a/util/process_wrapper/options.rs b/util/process_wrapper/options.rs index 24bba9fe80..5a074a76ff 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::ErrorFormat, } 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::ErrorFormat::Json), + "rendered" => Ok(rustc::ErrorFormat::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..049090c7fa --- /dev/null +++ b/util/process_wrapper/output.rs @@ -0,0 +1,56 @@ +// Copyright 2020 The Bazel Authors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::io::{self, prelude::*}; + +/// LineOutput tells process_output what to do when a line is processed. +/// If a Message is returned, it will be written to write_end, if +/// Skip is returned nothing will be printed and execution continues, +/// if Terminate is returned, process_output returns immediately. +/// Terminate is used to stop processing when we see an emit metadata +/// message. +#[derive(Debug)] +pub(crate) enum LineOutput { + Message(String), + Skip, + Terminate, +} + +/// process_output reads lines from read_end and invokes process_line on each. +/// Depending on the result of process_line, the modified message may be written +/// to write_end. +pub(crate) fn process_output( + read_end: Box, + write_end: Box, + mut process_line: F, +) -> io::Result<()> +where + F: FnMut(String) -> LineOutput, +{ + 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; + } + match process_line(line) { + LineOutput::Message(to_write) => writer.write_all(to_write.as_bytes())?, + LineOutput::Skip => {} + LineOutput::Terminate => return Ok(()), + }; + } + Ok(()) +} diff --git a/util/process_wrapper/rustc.rs b/util/process_wrapper/rustc.rs new file mode 100644 index 0000000000..e5667279b1 --- /dev/null +++ b/util/process_wrapper/rustc.rs @@ -0,0 +1,78 @@ +// Copyright 2020 The Bazel Authors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use tinyjson::JsonValue; + +use crate::output::LineOutput; + +#[derive(Debug, Copy, Clone)] +pub(crate) enum ErrorFormat { + Json, + Rendered, +} + +impl Default for ErrorFormat { + fn default() -> Self { + Self::Rendered + } +} + +fn get_key(value: &JsonValue, key: &str) -> Option { + if let JsonValue::Object(map) = value { + if let JsonValue::String(s) = map.get(key)? { + Some(s.clone()) + } else { + None + } + } else { + None + } +} + +/// stop_on_rmeta_completion takes an output line from rustc configured with +/// --error-format=json, parses the json and returns the appropriate output +/// according to the original --error-format supplied to rustc. +/// In addition, it will signal to stop when metadata is emitted +/// so the compiler can be terminated. +/// This is used to implement pipelining in rules_rust, please see +/// https://internals.rust-lang.org/t/evaluating-pipelined-rustc-compilation/10199 +pub(crate) fn stop_on_rmeta_completion( + line: String, + error_format: ErrorFormat, + kill: &mut bool, +) -> LineOutput { + let parsed: JsonValue = line + .parse() + .expect("process wrapper error: expected json messages in pipeline mode"); + if let Some(emit) = get_key(&parsed, "emit") { + // We don't want to print emit messages. + // If the emit messages is "metadata" we can signal the process to quit + return if emit == "metadata" { + *kill = true; + LineOutput::Terminate + } else { + LineOutput::Skip + }; + }; + + match error_format { + // If the output should be json, we just forward the messages as-is + ErrorFormat::Json => LineOutput::Message(line), + // Otherwise we extract the "rendered" attribute. + // If we don't find it we skip the line. + _ => get_key(&parsed, "rendered") + .map(LineOutput::Message) + .unwrap_or(LineOutput::Skip), + } +}