diff --git a/rust/private/rust.bzl b/rust/private/rust.bzl index 66293bf766..3346d8ba9a 100644 --- a/rust/private/rust.bzl +++ b/rust/private/rust.bzl @@ -690,6 +690,9 @@ RUSTC_ATTRS = { "_extra_rustc_flags": attr.label( default = Label("//rust/settings:extra_rustc_flags"), ), + "_experimental_persistent_worker": attr.label( + default = Label("//rust/settings:experimental_persistent_worker"), + ), "_per_crate_rustc_flag": attr.label( default = Label("//rust/settings:per_crate_rustc_flag"), ), diff --git a/rust/private/rustc.bzl b/rust/private/rustc.bzl index 04818522c7..0c097d09ce 100644 --- a/rust/private/rustc.bzl +++ b/rust/private/rustc.bzl @@ -917,6 +917,12 @@ def has_location_expansion(values): return True return False +_WORKER_RESPONSE_FILE_ARG = "--rules-rust-response-file=" + +def _worker_request_arg(arg): + """Prevent Bazel from recursively expanding rustc response files in a WorkRequest.""" + return _WORKER_RESPONSE_FILE_ARG + arg[1:] if arg.startswith("@") else arg + def _args_map_bin_dir(file): """Extract `bazel-out//bin` from a File whose path lives in the configuration's bin directory. @@ -956,7 +962,8 @@ def construct_arguments( skip_expanding_rustc_env = False, require_explicit_unstable_features = False, error_format = None, - allowed_unstable_rust_features = None): + allowed_unstable_rust_features = None, + use_persistent_worker = False): """Builds an Args object containing common rustc flags Args: @@ -1028,6 +1035,7 @@ def construct_arguments( require_explicit_unstable_features (bool): Whether to require all unstable features to be explicitly opted in to using `-Zallow-features=...`. error_format (str, optional): Error format to pass to the `--error-format` command line argument. If set to None, uses the "_error_format" entry in `attr`. allowed_unstable_rust_features (list, optional): List of unstable Rust language features allowed for this target. + use_persistent_worker (bool): Whether Bazel will construct worker requests from the rustc param file. Returns: tuple: A tuple of the following items @@ -1457,6 +1465,8 @@ def construct_arguments( # require_explicit_unstable_features makes no sense when all features are allowed anyway if require_explicit_unstable_features: process_wrapper_flags.add("--require-explicit-unstable-features", "true") + if use_persistent_worker: + extra_rustc_flags = [_worker_request_arg(flag) for flag in extra_rustc_flags] rustc_flags.add_all(extra_rustc_flags, map_each = map_flag) if is_no_std(ctx, toolchain, crate_info.is_test): @@ -1486,26 +1496,32 @@ def construct_arguments( elif rust_flags: for flag in _extract_allowed_unstable_features_from_flags(rust_flags, all_allowed_unstable_features): if type(flag) in ["tuple", "list"] and len(flag) == 2: + flag_format = _worker_request_arg(flag[0]) if use_persistent_worker else flag[0] rustc_flags.add_all( [flag[1]], - format_each = flag[0], + format_each = flag_format, expand_directories = False, ) else: if map_flag: flag = map_flag(flag) if flag != None: + if use_persistent_worker: + flag = _worker_request_arg(flag) rustc_flags.add(flag) # Add target specific flags last, so they can override previous flags authored_rustc_flags = getattr(attr, "rustc_flags", []) + expanded_authored_rustc_flags = expand_list_element_locations( + ctx, + authored_rustc_flags, + data_paths, + {}, + ) + if use_persistent_worker: + expanded_authored_rustc_flags = [_worker_request_arg(flag) for flag in expanded_authored_rustc_flags] rustc_flags.add_all( - expand_list_element_locations( - ctx, - authored_rustc_flags, - data_paths, - {}, - ), + expanded_authored_rustc_flags, map_each = map_flag, ) @@ -1713,6 +1729,10 @@ def rustc_compile_action( # # Exclude lints if we're building in the exec configuration to prevent crates # used in build scripts from generating warnings. + opaque_rust_flags = rust_flags if type(rust_flags) == "Args" else None + if opaque_rust_flags != None: + rust_flags = [] + lint_files = [] if hasattr(ctx.attr, "lint_config") and ctx.attr.lint_config and not is_exec_configuration(ctx): rust_flags = rust_flags + ctx.attr.lint_config[LintsInfo].rustc_lint_flags @@ -1794,6 +1814,14 @@ def rustc_compile_action( dwo_outputs = ctx.actions.declare_directory(fission_directory, sibling = crate_info.output) rust_flags.append(("-Zsplit-dwarf-out-dir=%s", dwo_outputs)) + use_persistent_worker = ( + bool(ctx.executable._process_wrapper) and + hasattr(ctx.attr, "_experimental_persistent_worker") and + ctx.attr._experimental_persistent_worker[BuildSettingInfo].value and + opaque_rust_flags == None + ) + use_incremental = use_persistent_worker and ctx.var["COMPILATION_MODE"] != "opt" and not _will_emit_object_file(emit) + args, env_from_args = construct_arguments( ctx = ctx, attr = attr, @@ -1818,7 +1846,12 @@ def rustc_compile_action( skip_expanding_rustc_env = skip_expanding_rustc_env, require_explicit_unstable_features = require_explicit_unstable_features, allowed_unstable_rust_features = allowed_unstable_rust_features, + use_persistent_worker = use_persistent_worker, ) + if opaque_rust_flags != None: + args.all.append(opaque_rust_flags) + if use_incremental: + args.process_wrapper_flags.add("--rustc-incremental", "true") args_metadata = None if build_metadata: @@ -1846,7 +1879,10 @@ def rustc_compile_action( build_metadata = True, require_explicit_unstable_features = require_explicit_unstable_features, allowed_unstable_rust_features = allowed_unstable_rust_features, + use_persistent_worker = use_persistent_worker, ) + if opaque_rust_flags != None: + args_metadata.all.append(opaque_rust_flags) env = dict(ctx.configuration.default_shell_env) @@ -1893,6 +1929,15 @@ def rustc_compile_action( action_outputs.append(dwo_outputs) # buildifier: disable=uninitialized if ctx.executable._process_wrapper: + execution_requirements = {} + if args.supports_path_mapping: + execution_requirements["supports-path-mapping"] = "" + if use_persistent_worker: + execution_requirements.update({ + "requires-worker-protocol": "proto", + "supports-workers": "1", + }) + # Run as normal ctx.actions.run( executable = ctx.executable._process_wrapper, @@ -1910,7 +1955,7 @@ def rustc_compile_action( ), toolchain = "@rules_rust//rust:toolchain_type", resource_set = get_rustc_resource_set(toolchain), - execution_requirements = {"supports-path-mapping": ""} if args.supports_path_mapping else None, + execution_requirements = execution_requirements, ) if args_metadata: ctx.actions.run( @@ -1928,7 +1973,7 @@ def rustc_compile_action( "" if len(srcs) == 1 else "s", ), toolchain = "@rules_rust//rust:toolchain_type", - execution_requirements = {"supports-path-mapping": ""} if args_metadata.supports_path_mapping else None, + execution_requirements = execution_requirements, ) elif hasattr(ctx.executable, "_bootstrap_process_wrapper"): # Run without process_wrapper @@ -1939,7 +1984,7 @@ def rustc_compile_action( inputs = compile_inputs, outputs = action_outputs, env = env, - arguments = [args.rustc_path, args.rustc_flags], + arguments = [args.rustc_path, args.rustc_flags] + ([opaque_rust_flags] if opaque_rust_flags != None else []), mnemonic = "Rustc", progress_message = "Compiling Rust (without process_wrapper) {} {}{} ({} file{})".format( crate_info.type, diff --git a/rust/settings/BUILD.bazel b/rust/settings/BUILD.bazel index 454afd8de2..222bea3815 100644 --- a/rust/settings/BUILD.bazel +++ b/rust/settings/BUILD.bazel @@ -14,6 +14,7 @@ load( "error_format", "experimental_compile_rustdoc_tests", "experimental_link_std_dylib", + "experimental_persistent_worker", "experimental_use_allocator_libraries_with_mangled_symbols", "experimental_use_cc_common_link", "experimental_use_coverage_metadata_files", @@ -87,6 +88,8 @@ experimental_compile_rustdoc_tests() experimental_link_std_dylib() +experimental_persistent_worker() + experimental_use_cc_common_link() experimental_use_coverage_metadata_files() diff --git a/rust/settings/settings.bzl b/rust/settings/settings.bzl index a25281e10e..b67f14521d 100644 --- a/rust/settings/settings.bzl +++ b/rust/settings/settings.bzl @@ -115,6 +115,17 @@ def pipelined_compilation(): build_setting_default = False, ) +def experimental_persistent_worker(): + """Enable persistent rustc workers and incremental compilation for non-opt builds. + + Actions remain compatible with one-shot local and remote execution. Persistent-worker-capable + executors can reuse rustc's incremental cache between requests. + """ + bool_flag( + name = "experimental_persistent_worker", + build_setting_default = False, + ) + # buildifier: disable=unnamed-macro def experimental_use_cc_common_link(): """A flag to control whether to link rust_binary and rust_test targets using \ diff --git a/util/process_wrapper/main.rs b/util/process_wrapper/main.rs index 39a6d6db16..cefae98922 100644 --- a/util/process_wrapper/main.rs +++ b/util/process_wrapper/main.rs @@ -17,6 +17,7 @@ mod options; mod output; mod rustc; mod util; +mod worker; use std::collections::HashMap; use std::fmt; @@ -118,11 +119,29 @@ fn process_line( } fn main() -> Result<(), ProcessWrapperError> { + let argv: Vec = std::env::args().collect(); + if let Some(worker_arg_index) = argv.iter().position(|arg| arg == "--persistent_worker") { + let mut startup_argv = argv; + startup_argv.truncate(worker_arg_index); + return worker::run(startup_argv).map_err(|e| ProcessWrapperError(e.to_string())); + } + let opts = options().map_err(|e| ProcessWrapperError(e.to_string()))?; + let incremental_cache = if opts.rustc_incremental { + opts.rustc_incremental_dir + } else { + None + }; + let mut command = Command::new(opts.executable); command .args(opts.child_arguments) + .args( + incremental_cache + .as_ref() + .map(|cache| format!("-Cincremental={cache}")), + ) .env_clear() .envs(opts.child_environment) .stdout(if let Some(stdout_file) = opts.stdout_file { diff --git a/util/process_wrapper/options.rs b/util/process_wrapper/options.rs index cca227bb0a..ea82fa6f25 100644 --- a/util/process_wrapper/options.rs +++ b/util/process_wrapper/options.rs @@ -9,6 +9,8 @@ use crate::flags::{FlagParseError, Flags, ParseOutcome}; use crate::rustc; use crate::util::*; +const WORKER_RESPONSE_FILE_ARG: &str = "--rules-rust-response-file="; + #[derive(Debug)] pub(crate) enum OptionError { FlagError(FlagParseError), @@ -49,6 +51,10 @@ pub(crate) struct Options { pub(crate) rustc_quit_on_rmeta: bool, // This controls the output format of rustc messages. pub(crate) rustc_output_format: Option, + // Enable rustc incremental compilation using a cache owned by this wrapper. + pub(crate) rustc_incremental: bool, + // An optional persistent cache directory supplied by the worker parent. + pub(crate) rustc_incremental_dir: Option, } pub(crate) fn options() -> Result { @@ -67,6 +73,8 @@ pub(crate) fn options() -> Result { let mut output_file = None; let mut rustc_quit_on_rmeta_raw = None; let mut rustc_output_format_raw = None; + let mut rustc_incremental_raw = None; + let mut rustc_incremental_dir = None; let mut flags = Flags::new(); let mut require_explicit_unstable_features = None; flags.define_repeated_flag("--subst", "", &mut subst_mapping_raw); @@ -126,6 +134,16 @@ pub(crate) fn options() -> Result { Default: `rendered`", &mut rustc_output_format_raw, ); + flags.define_flag( + "--rustc-incremental", + "Enable rustc incremental compilation.", + &mut rustc_incremental_raw, + ); + flags.define_flag( + "--rustc-incremental-dir", + "Persistent incremental cache directory supplied by the worker.", + &mut rustc_incremental_dir, + ); flags.define_flag( "--require-explicit-unstable-features", "If set, an empty -Zallow-features= will be added to the rustc command line whenever no \ @@ -252,6 +270,7 @@ pub(crate) fn options() -> Result { .transpose()?; let rustc_quit_on_rmeta = rustc_quit_on_rmeta_raw.is_some_and(|s| s == "true"); + let rustc_incremental = rustc_incremental_raw.is_some_and(|s| s == "true"); let rustc_output_format = rustc_output_format_raw .map(|v| match v.as_str() { "json" => Ok(rustc::ErrorFormat::Json), @@ -301,6 +320,8 @@ pub(crate) fn options() -> Result { output_file, rustc_quit_on_rmeta, rustc_output_format, + rustc_incremental, + rustc_incremental_dir, }) } @@ -344,6 +365,11 @@ fn prepare_arg(mut arg: String, subst_mappings: &[(String, String)]) -> String { arg } +fn response_file_path(arg: &str) -> Option<&str> { + arg.strip_prefix(WORKER_RESPONSE_FILE_ARG) + .or_else(|| arg.strip_prefix('@')) +} + /// Apply substitutions to the given param file. Returns true iff any allow-features flags were found. fn prepare_param_file( filename: &str, @@ -361,7 +387,7 @@ fn prepare_param_file( for arg in read_file(filename)? { let arg = prepare_arg(arg, subst_mappings); has_allow_features_flag |= is_allow_features_flag(&arg); - if let Some(arg_file) = arg.strip_prefix('@') { + if let Some(arg_file) = response_file_path(&arg) { has_allow_features_flag |= process_file(arg_file, subst_mappings, read_file, write_to_file)?; } else { @@ -391,7 +417,7 @@ fn prepare_args( for arg in args.into_iter() { let arg = prepare_arg(arg, subst_mappings); - if let Some(param_file) = arg.strip_prefix('@') { + if let Some(param_file) = response_file_path(&arg) { let expanded_file = format!("{param_file}.expanded"); let format_err = |err: io::Error| { OptionError::Generic(format!( @@ -474,6 +500,47 @@ fn environment_block( mod test { use super::*; + #[test] + fn test_worker_response_file_applies_substitutions() { + let mut written_files = HashMap::::new(); + let mut read_file = |filename: &str| -> Result, OptionError> { + match filename { + "/exec/flags.params" => Ok(vec!["@${pwd}/nested.params".to_owned()]), + "/exec/nested.params" => Ok(vec!["--cfg=from_response_file".to_owned()]), + _ => Err(OptionError::Generic(format!( + "unexpected response file: {filename}" + ))), + } + }; + let mut write_file = |filename: &str, content: &str| -> Result<(), OptionError> { + written_files.insert(filename.to_owned(), content.to_owned()); + Ok(()) + }; + + let args = vec![ + "rustc".to_owned(), + format!("{WORKER_RESPONSE_FILE_ARG}${{pwd}}/flags.params"), + ]; + let subst_mappings = vec![("pwd".to_owned(), "/exec".to_owned())]; + let args = prepare_args( + args, + &subst_mappings, + false, + Some(&mut read_file), + Some(&mut write_file), + ) + .unwrap(); + + assert_eq!(args, ["rustc", "@/exec/flags.params.expanded"]); + assert_eq!( + written_files, + HashMap::from([( + "/exec/flags.params.expanded".to_owned(), + "--cfg=from_response_file".to_owned() + )]) + ); + } + #[test] fn test_enforce_allow_features_flag_user_didnt_say() { let args = vec!["rustc".to_string()]; diff --git a/util/process_wrapper/worker.rs b/util/process_wrapper/worker.rs new file mode 100644 index 0000000000..9cb6a21e2c --- /dev/null +++ b/util/process_wrapper/worker.rs @@ -0,0 +1,350 @@ +// Copyright 2026 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::collections::hash_map::DefaultHasher; +use std::convert::TryFrom; +use std::fs; +use std::hash::{Hash, Hasher}; +use std::io::{self, BufReader, BufWriter, Read, Write}; +use std::path::{Path, PathBuf}; +use std::process::Command; + +const MAX_REQUEST_SIZE: usize = 64 * 1024 * 1024; + +struct WorkerDir(PathBuf); + +impl Drop for WorkerDir { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } +} + +fn create_worker_dir() -> io::Result { + let path = std::env::temp_dir() + .join("rules_rust_worker") + .join(std::process::id().to_string()); + let _ = fs::remove_dir_all(&path); + fs::create_dir_all(&path)?; + Ok(WorkerDir(path)) +} + +#[derive(Debug, Default, PartialEq)] +struct WorkRequest { + arguments: Vec, + request_id: i32, + cancel: bool, +} + +#[derive(Debug, Default, PartialEq)] +struct WorkResponse { + exit_code: i32, + output: String, + request_id: i32, + was_cancelled: bool, +} + +pub(crate) fn run(startup_argv: Vec) -> io::Result<()> { + let incremental = startup_argv + .windows(2) + .any(|args| args == ["--rustc-incremental", "true"]); + let worker_dir = create_worker_dir()?; + let cache = if incremental { + let path = worker_dir.0.join("incremental"); + fs::create_dir_all(&path)?; + Some(path) + } else { + None + }; + + let stdin = io::stdin(); + let stdout = io::stdout(); + let mut reader = BufReader::new(stdin.lock()); + let mut writer = BufWriter::new(stdout.lock()); + + while let Some(request) = read_request(&mut reader)? { + let response = if request.cancel { + WorkResponse { + request_id: request.request_id, + was_cancelled: true, + ..WorkResponse::default() + } + } else { + execute_request(&startup_argv, request, &worker_dir.0, cache.as_deref()) + }; + write_response(&mut writer, &response)?; + writer.flush()?; + } + Ok(()) +} + +fn execute_request( + startup_argv: &[String], + request: WorkRequest, + worker_dir: &Path, + cache: Option<&Path>, +) -> WorkResponse { + let request_id = request.request_id; + let result = (|| { + let request_cache = cache + .map(|cache_root| { + let mut hasher = DefaultHasher::new(); + request.arguments.hash(&mut hasher); + let path = cache_root.join(format!("{:016x}", hasher.finish())); + fs::create_dir_all(&path)?; + Ok::<_, io::Error>(path) + }) + .transpose()?; + let request_file = worker_dir.join("request.params"); + let mut writer = BufWriter::new(fs::File::create(&request_file)?); + for arg in &request.arguments { + writeln!(writer, "{arg}")?; + } + writer.flush()?; + drop(writer); + + let mut argv = startup_argv[1..].to_vec(); + let delimiter = argv + .iter() + .position(|arg| arg == "--") + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "missing -- delimiter"))?; + if let Some(cache) = request_cache { + argv.splice( + delimiter..delimiter, + [ + "--rustc-incremental-dir".to_owned(), + cache.to_string_lossy().into_owned(), + ], + ); + } + Command::new(&startup_argv[0]) + .args(argv) + .arg(format!("@{}", request_file.display())) + .output() + })(); + + match result { + Ok(output) => { + let mut combined = output.stdout; + combined.extend(output.stderr); + WorkResponse { + exit_code: output.status.code().unwrap_or(1), + output: String::from_utf8_lossy(&combined).into_owned(), + request_id, + was_cancelled: false, + } + } + Err(error) => WorkResponse { + exit_code: 1, + output: format!("process wrapper worker failed to execute request: {error}\n"), + request_id, + was_cancelled: false, + }, + } +} + +fn read_request(reader: &mut impl Read) -> io::Result> { + let Some(length) = read_varint(reader)? else { + return Ok(None); + }; + let length = usize::try_from(length) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "request is too large"))?; + if length > MAX_REQUEST_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "request exceeds maximum size", + )); + } + let mut bytes = vec![0; length]; + reader.read_exact(&mut bytes)?; + decode_proto_request(&bytes).map(Some) +} + +fn decode_proto_request(bytes: &[u8]) -> io::Result { + let mut request = WorkRequest::default(); + let mut offset = 0; + while offset < bytes.len() { + let key = read_varint_from_slice(bytes, &mut offset)?; + let field = key >> 3; + let wire_type = key & 7; + match (field, wire_type) { + (1, 2) => { + let value = read_length_delimited(bytes, &mut offset)?; + request.arguments.push( + std::str::from_utf8(value) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))? + .to_owned(), + ); + } + (3, 0) => request.request_id = read_varint_from_slice(bytes, &mut offset)? as i32, + (4, 0) => request.cancel = read_varint_from_slice(bytes, &mut offset)? != 0, + _ => skip_proto_field(bytes, &mut offset, wire_type)?, + } + } + Ok(request) +} + +fn write_response(writer: &mut impl Write, response: &WorkResponse) -> io::Result<()> { + let mut bytes = Vec::new(); + if response.exit_code != 0 { + write_proto_key(&mut bytes, 1, 0); + write_varint(&mut bytes, response.exit_code as u32 as u64)?; + } + if !response.output.is_empty() { + write_proto_key(&mut bytes, 2, 2); + write_varint(&mut bytes, response.output.len() as u64)?; + bytes.extend(response.output.as_bytes()); + } + if response.request_id != 0 { + write_proto_key(&mut bytes, 3, 0); + write_varint(&mut bytes, response.request_id as u32 as u64)?; + } + if response.was_cancelled { + write_proto_key(&mut bytes, 4, 0); + write_varint(&mut bytes, 1)?; + } + write_varint(writer, bytes.len() as u64)?; + writer.write_all(&bytes) +} + +fn read_varint(reader: &mut impl Read) -> io::Result> { + let mut value = 0u64; + for shift in (0..70).step_by(7) { + let mut byte = [0]; + match reader.read_exact(&mut byte) { + Ok(()) => {} + Err(error) if error.kind() == io::ErrorKind::UnexpectedEof && shift == 0 => { + return Ok(None) + } + Err(error) => return Err(error), + } + if shift == 63 && byte[0] > 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "protobuf varint overflow", + )); + } + value |= u64::from(byte[0] & 0x7f) << shift; + if byte[0] & 0x80 == 0 { + return Ok(Some(value)); + } + } + Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid protobuf varint", + )) +} + +fn read_varint_from_slice(bytes: &[u8], offset: &mut usize) -> io::Result { + let mut value = 0u64; + for shift in (0..70).step_by(7) { + let byte = *bytes + .get(*offset) + .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "truncated varint"))?; + *offset += 1; + if shift == 63 && byte > 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "protobuf varint overflow", + )); + } + value |= u64::from(byte & 0x7f) << shift; + if byte & 0x80 == 0 { + return Ok(value); + } + } + Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid protobuf varint", + )) +} + +fn read_length_delimited<'a>(bytes: &'a [u8], offset: &mut usize) -> io::Result<&'a [u8]> { + let length = usize::try_from(read_varint_from_slice(bytes, offset)?) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "field is too large"))?; + let end = offset + .checked_add(length) + .filter(|end| *end <= bytes.len()) + .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "truncated field"))?; + let value = &bytes[*offset..end]; + *offset = end; + Ok(value) +} + +fn skip_proto_field(bytes: &[u8], offset: &mut usize, wire_type: u64) -> io::Result<()> { + let length = match wire_type { + 0 => { + read_varint_from_slice(bytes, offset)?; + return Ok(()); + } + 1 => 8, + 2 => usize::try_from(read_varint_from_slice(bytes, offset)?) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "field is too large"))?, + 5 => 4, + _ => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "unsupported protobuf wire type", + )) + } + }; + *offset = offset + .checked_add(length) + .filter(|end| *end <= bytes.len()) + .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "truncated field"))?; + Ok(()) +} + +fn write_proto_key(bytes: &mut Vec, field: u64, wire_type: u64) { + write_varint(bytes, (field << 3) | wire_type).expect("writing to a Vec cannot fail"); +} + +fn write_varint(writer: &mut impl Write, mut value: u64) -> io::Result<()> { + loop { + let mut byte = (value & 0x7f) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + writer.write_all(&[byte])?; + if value == 0 { + return Ok(()); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn decodes_and_encodes_work_messages() { + let request = [ + 11, 0x0a, 0x03, b'f', b'o', b'o', 0x12, 0x02, 0x08, 0x01, 0x18, 0x07, + ]; + let mut request = request.as_slice(); + let request = read_request(&mut request).unwrap().unwrap(); + assert_eq!(request.arguments, ["foo"]); + assert_eq!(request.request_id, 7); + + let response = WorkResponse { + exit_code: 2, + output: "no".to_owned(), + request_id: 7, + was_cancelled: false, + }; + let mut bytes = Vec::new(); + write_response(&mut bytes, &response).unwrap(); + assert_eq!(bytes, [8, 0x08, 0x02, 0x12, 0x02, b'n', b'o', 0x18, 0x07]); + } +}