diff --git a/.cargo/config.toml b/.cargo/config.toml index 70397b4..a4ae00c 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,3 +1,3 @@ [profile.profiling] inherits = "release" -debug = false +debug = true diff --git a/Cargo.lock b/Cargo.lock index e87d86d..fe2e7d6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -968,6 +968,7 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" name = "lib" version = "0.6.0" dependencies = [ + "aho-corasick", "anyhow", "base64", "derive_more", diff --git a/Cargo.toml b/Cargo.toml index 616d1d9..ae1b610 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,6 @@ futures-util = { version = "0.3.31", features = ["sink", "std"] } clap = { version = "4.5.50", features = ["derive", "string", "cargo"] } clap-verbosity-flag = "3.0.4" serde = { version = "1.0.228", features = ["derive", "rc"] } -serde_json = { version = "1.0.145" } tokio = { version = "1.48.0", features = ["full"] } tracing = { version = "0.1.41", features = ["release_max_level_debug"] } tracing-log = "0.2.0" @@ -30,3 +29,12 @@ nix-compat = { git = "https://git.snix.dev/snix/snix.git", features = [ "serde", "flakeref", ] } +# simd-json = { version = "0.17.0", features = [ +# "serde_impl", +# # swar-number-parsing is disabled because nix never outputs +# # floats. +# # "swar-number-parsing", +# "runtime-detection", +# "hints", +# ] } +serde_json = { version = "1.0.145" } diff --git a/nix/shells.nix b/nix/shells.nix index 24177db..d291524 100644 --- a/nix/shells.nix +++ b/nix/shells.nix @@ -21,6 +21,7 @@ pkgs.just pkgs.pnpm pkgs.nodejs + pkgs.perf ]; PROTOC = lib.getExe pkgs.protobuf; diff --git a/wire/cli/default.nix b/wire/cli/default.nix index 260bdfa..0d4a4ec 100644 --- a/wire/cli/default.nix +++ b/wire/cli/default.nix @@ -23,7 +23,9 @@ pname = "wire"; cargoExtraArgs = "-p wire"; doCheck = true; - nativeBuildInputs = [ pkgs.installShellFiles ]; + nativeBuildInputs = [ + pkgs.installShellFiles + ]; postInstall = '' installShellCompletion --cmd wire \ --bash <($out/bin/wire completions bash) \ @@ -36,6 +38,12 @@ CARGO_PROFILE = "dev"; }; + wire-unwrapped-perf = buildRustProgram { + name = "wire"; + pname = "wire"; + cargoExtraArgs = "-p wire --features dhat-heap"; + }; + wire = pkgs.symlinkJoin { name = "wire"; paths = [ self'.packages.wire-unwrapped ]; @@ -68,6 +76,10 @@ paths = [ self'.packages.wire-unwrapped-dev ]; }; + wire-small-perf = self'.packages.wire-small.overrideAttrs { + paths = [ self'.packages.wire-unwrapped-perf ]; + }; + wire-dignostics-md = self'.packages.wire-unwrapped.overrideAttrs { DIAGNOSTICS_MD_OUTPUT = "/build/source"; installPhase = '' diff --git a/wire/lib/Cargo.toml b/wire/lib/Cargo.toml index 2724981..47adb41 100644 --- a/wire/lib/Cargo.toml +++ b/wire/lib/Cargo.toml @@ -33,6 +33,7 @@ sha2 = { workspace = true } base64 = { workspace = true } nix-compat = { workspace = true } strip-ansi-escapes = "0.2.1" +aho-corasick = "1.1.3" [dev-dependencies] tempdir = "0.3" diff --git a/wire/lib/src/commands/interactive.rs b/wire/lib/src/commands/interactive.rs index 8aee18d..3016923 100644 --- a/wire/lib/src/commands/interactive.rs +++ b/wire/lib/src/commands/interactive.rs @@ -62,9 +62,9 @@ struct CompletionStatus { struct WatchStdoutArguments { began_tx: Sender<()>, reader: MasterReader, - succeed_needle: Arc, - failed_needle: Arc, - start_needle: Arc, + succeed_needle: Arc>, + failed_needle: Arc>, + start_needle: Arc>, output_mode: ChildOutputMode, stderr_collection: Arc>>, stdout_collection: Arc>>, @@ -77,6 +77,8 @@ struct WatchStdoutArguments { const THREAD_BEGAN_SIGNAL: &[u8; 1] = b"b"; const THREAD_QUIT_SIGNAL: &[u8; 1] = b"q"; +const NEEDLE_LENGTH: usize = 9; + /// substitutes STDOUT with #$line. stdout is far less common than stderr. const IO_SUBS: &str = "1> >(while IFS= read -r line; do echo \"#$line\"; done)"; @@ -100,9 +102,9 @@ pub(crate) fn interactive_command_with_env>( let command_string = &format!( "echo '{start}' && {command} {flags} {IO_SUBS} && echo '{succeed}' || echo '{failed}'", - start = start_needle, - succeed = succeed_needle, - failed = failed_needle, + start = String::from_utf8_lossy(&start_needle), + succeed = String::from_utf8_lossy(&succeed_needle), + failed = String::from_utf8_lossy(&failed_needle), command = arguments.command_string.as_ref(), flags = match arguments.output_mode { ChildOutputMode::Nix => "--log-format internal-json", @@ -208,13 +210,13 @@ pub(crate) fn interactive_command_with_env>( }) } -fn create_needles() -> (Arc, Arc, Arc) { +fn create_needles() -> (Arc>, Arc>, Arc>) { let tmp_prefix = rand::distr::SampleString::sample_string(&Alphabetic, &mut rand::rng(), 5); ( - Arc::new(format!("{tmp_prefix}_WIRE_QUIT")), - Arc::new(format!("{tmp_prefix}_WIRE_FAIL")), - Arc::new(format!("{tmp_prefix}_WIRE_START")), + Arc::new(format!("{tmp_prefix}_W_Q").as_bytes().to_vec()), + Arc::new(format!("{tmp_prefix}_W_F").as_bytes().to_vec()), + Arc::new(format!("{tmp_prefix}_W_S").as_bytes().to_vec()), ) } @@ -420,41 +422,45 @@ fn dynamic_watch_sudo_stdout(arguments: WatchStdoutArguments) -> Result<(), Comm match reader.read(&mut buffer) { Ok(0) => break 'outer, Ok(n) => { - let new_data = String::from_utf8_lossy(&buffer[..n]); - log_buffer.process(&new_data); + log_buffer.process_slice(&buffer[..n]); + + while let Some(mut line) = log_buffer.next_line() { + let mut windows = line.windows(NEEDLE_LENGTH); - for line in log_buffer.take_lines() { - if line.contains(start_needle.as_ref()) { - debug!("{start_needle} was found, switching mode..."); + if windows.any(|window| window == *start_needle) { + debug!("start needle was found, switching mode..."); let _ = began_tx.send(()); began = true; continue; } - if line.contains(succeed_needle.as_ref()) { - debug!("{succeed_needle} was found, marking child as succeeding."); + if windows.any(|window| window == *succeed_needle) { + debug!("succeed needle was found, marking child as succeeding."); completion_status.mark_completed(true); break 'outer; } - if line.contains(failed_needle.as_ref()) { - debug!("{failed_needle} was found, elevated child did not succeed."); + if windows.any(|window| window == *failed_needle) { + debug!("failed needle was found, elevated child did not succeed."); completion_status.mark_completed(false); break 'outer; } if began { - if let Some(stripped) = line.strip_prefix('#') { + if line.starts_with(b"#") { + let stripped = &mut line[1..]; + if log_stdout { - output_mode.trace(&stripped.to_string()); + output_mode.trace_slice(stripped); } let mut queue = stdout_collection.lock().unwrap(); - queue.push_front(stripped.to_string()); + // clone + queue.push_front(String::from_utf8_lossy(stripped).to_string()); continue; } - let log = output_mode.trace(&line); + let log = output_mode.trace_slice(&mut line); let mut queue = stderr_collection.lock().unwrap(); if let SubcommandLog::Internal(log) = log { @@ -466,7 +472,7 @@ fn dynamic_watch_sudo_stdout(arguments: WatchStdoutArguments) -> Result<(), Comm } } else { stdout - .write_all(new_data.as_bytes()) + .write_all(&line) .map_err(CommandError::WritingClientStdout)?; stdout.flush().map_err(CommandError::WritingClientStdout)?; } diff --git a/wire/lib/src/commands/interactive_logbuffer.rs b/wire/lib/src/commands/interactive_logbuffer.rs index eac9692..40089a5 100644 --- a/wire/lib/src/commands/interactive_logbuffer.rs +++ b/wire/lib/src/commands/interactive_logbuffer.rs @@ -1,32 +1,37 @@ -use std::collections::VecDeque; +// SPDX-License-Identifier: AGPL-3.0-or-later +// Copyright 2024-2025 wire Contributors /// Split into its own struct to be tested nicer pub(crate) struct LogBuffer { - buffer: String, - lines: VecDeque, + buffer: Vec, } impl LogBuffer { pub fn new() -> Self { - Self { - buffer: String::new(), - lines: VecDeque::new(), - } + Self { buffer: Vec::new() } + } + + pub fn process_slice(&mut self, slice: &[u8]) { + self.buffer.extend_from_slice(slice); } - pub fn process(&mut self, new_data: &str) { - self.buffer.push_str(new_data); + pub fn next_line(&mut self) -> Option> { + let line_end = self.buffer.iter().position(|x| *x == b'\n')?; - while let Some(newline) = self.buffer.find('\n') { - let line = self.buffer[..newline].to_string(); - self.buffer = self.buffer[newline + 1..].to_string(); - self.lines.push_back(line); - } + let drained = self.buffer.drain(..line_end).collect(); + self.buffer.remove(0); + Some(drained) } - /// deletes old lines and gives the current ones. - pub fn take_lines(&mut self) -> VecDeque { - std::mem::take(&mut self.lines) + #[cfg(test)] + fn take_lines(&mut self) -> Vec> { + let mut lines = vec![]; + + while let Some(line) = self.next_line() { + lines.push(line); + } + + lines } } @@ -38,21 +43,21 @@ mod tests { fn test_split_line_processing() { let mut log_buffer = LogBuffer::new(); - log_buffer.process("Writing key KeySpec { destination: \"/et"); - log_buffer.process("c/keys/buildbot.aws.key\", user: \"buildbot\", group: \"buildbot-worker\", permissions: 384, length: 32, last: false, crc: 1370815231 }, 32 bytes of data"); - log_buffer.process("\n"); - log_buffer.process("xxx"); - log_buffer.process("xx_WIRE"); - log_buffer.process("_QUIT\n"); + log_buffer.process_slice(b"Writing key KeySpec { destination: \"/et"); + log_buffer.process_slice(b"c/keys/buildbot.aws.key\", user: \"buildbot\", group: \"buildbot-worker\", permissions: 384, length: 32, last: false, crc: 1370815231 }, 32 bytes of data"); + log_buffer.process_slice(b"\n"); + log_buffer.process_slice(b"xxx"); + log_buffer.process_slice(b"xx_WIRE"); + log_buffer.process_slice(b"_QUIT\n"); let lines = log_buffer.take_lines(); assert_eq!(lines.len(), 2); assert_eq!( - lines.front().unwrap(), + String::from_utf8_lossy(lines.first().unwrap()), "Writing key KeySpec { destination: \"/etc/keys/buildbot.aws.key\", user: \"buildbot\", group: \"buildbot-worker\", permissions: 384, length: 32, last: false, crc: 1370815231 }, 32 bytes of data" ); - assert_eq!(lines.get(1).unwrap(), "xxxxx_WIRE_QUIT"); + assert_eq!(lines.get(1), Some(&"xxxxx_WIRE_QUIT".as_bytes().to_vec())); // taking leaves none - assert_eq!(log_buffer.lines.len(), 0); + assert_eq!(log_buffer.take_lines().len(), 0); } } diff --git a/wire/lib/src/commands/mod.rs b/wire/lib/src/commands/mod.rs index 0e06e50..8d8f533 100644 --- a/wire/lib/src/commands/mod.rs +++ b/wire/lib/src/commands/mod.rs @@ -3,11 +3,11 @@ use std::{ collections::HashMap, - sync::{Arc, Mutex}, + sync::{Arc, LazyLock, Mutex}, }; +use aho_corasick::AhoCorasick; use nix_compat::log::{AT_NIX_PREFIX, LogMessage}; -use tracing::debug; use crate::{ SubCommandModifiers, @@ -49,6 +49,14 @@ pub(crate) struct CommandArguments<'t, S: AsRef> { log_stdout: bool, } +static AHO_CORASICK: LazyLock = LazyLock::new(|| { + AhoCorasick::builder() + .ascii_case_insensitive(false) + .match_kind(aho_corasick::MatchKind::LeftmostFirst) + .build([AT_NIX_PREFIX]) + .unwrap() +}); + impl<'a, S: AsRef> CommandArguments<'a, S> { pub(crate) fn new( command_string: S, @@ -142,26 +150,27 @@ impl WireCommandChip for Either { } impl ChildOutputMode { - fn trace(self, line: &String) -> nix_log::SubcommandLog<'_> { + /// this function is by far the biggest hotspot in the whole tree + fn trace_slice(self, line: &mut [u8]) -> nix_log::SubcommandLog<'_> { let log = match self { - ChildOutputMode::Nix => { - let stripped = line - .find(AT_NIX_PREFIX) - .map(|position| &line[position + AT_NIX_PREFIX.len()..]); - - if let Some(line) = stripped { - serde_json::from_str::(line).map_or_else( - |err| { - debug!("failed to parse {line:?}: {err:?}"); - SubcommandLog::Raw(line.into()) - }, - SubcommandLog::Internal, - ) - } else { - SubcommandLog::Raw(line.into()) + Self::Raw => SubcommandLog::Raw(String::from_utf8_lossy(line).to_string()), + Self::Nix => { + let line = AHO_CORASICK.find(&line).map(|x| &mut line[x.end()..]); + + if let Some(line) = line { + let log = + serde_json::from_slice::(line).map(SubcommandLog::Internal); + + match log { + Ok(log) => return log, + Err(err) => { + return SubcommandLog::Raw(format!("parsing log failed: {err:?}")); + } + } } + + SubcommandLog::Raw("line did not have a needle".to_string()) } - Self::Raw => SubcommandLog::Raw(line.into()), }; log.trace(); diff --git a/wire/lib/src/commands/noninteractive.rs b/wire/lib/src/commands/noninteractive.rs index 0026d41..84f5eaf 100644 --- a/wire/lib/src/commands/noninteractive.rs +++ b/wire/lib/src/commands/noninteractive.rs @@ -166,15 +166,17 @@ pub async fn handle_io( let mut io_reader = tokio::io::AsyncBufReadExt::lines(BufReader::new(reader)); while let Some(line) = io_reader.next_line().await.unwrap() { + let mut line = line.into_bytes(); + let log = if should_log { - Some(output_mode.trace(&line)) + Some(output_mode.trace_slice(&mut line)) } else { None }; if !is_error { let mut queue = collection.lock().await; - queue.push_front(line); + queue.push_front(String::from_utf8_lossy(&line).to_string()); } else if let Some(SubcommandLog::Internal(log)) = log { if let Some(message) = get_errorish_message(&log) { let mut queue = collection.lock().await; diff --git a/wire/lib/src/nix_log.rs b/wire/lib/src/nix_log.rs index 9142316..711038c 100644 --- a/wire/lib/src/nix_log.rs +++ b/wire/lib/src/nix_log.rs @@ -8,12 +8,10 @@ use std::{ }; use tracing::{Level as tracing_level, event, warn}; -// static DIGEST_RE: LazyLock = LazyLock::new(|| Regex::new(r"[0-9a-z]{32}").unwrap()); - #[derive(Debug)] pub enum SubcommandLog<'a> { Internal(LogMessage<'a>), - Raw(Cow<'a, str>), + Raw(String), } pub(crate) trait Trace { @@ -93,13 +91,14 @@ impl Trace for SubcommandLog<'_> { match self { SubcommandLog::Internal(line) => { line.trace(); + } + SubcommandLog::Raw(line) => { + if line.is_empty() { + return; + } - // tracing_indicatif::span_ext::IndicatifSpanExt::pb_set_message( - // &Span::current(), - // &DIGEST_RE.replace_all(&line.to_string(), "…"), - // ); + warn!("{line}"); } - SubcommandLog::Raw(line) => warn!("{line}"), } } }