diff --git a/crates/core/src/repository/warm_up.rs b/crates/core/src/repository/warm_up.rs index 467878f7..202e100f 100644 --- a/crates/core/src/repository/warm_up.rs +++ b/crates/core/src/repository/warm_up.rs @@ -1,12 +1,13 @@ -use std::io; -use std::process::Command; -use std::thread::sleep; +use std::io::{self, BufRead, BufReader}; +use std::process::{Command, Stdio}; +use std::thread::{self, sleep}; use backon::{BlockingRetryable, ExponentialBuilder}; use itertools::Itertools; -use log::{debug, error, warn}; +use log::{debug, error, info, warn}; use rayon::ThreadPoolBuilder; +use serde::Deserialize; use crate::{ CommandInput, Id, Progress, @@ -28,6 +29,70 @@ pub(super) mod constants { pub(crate) const INITIAL_DELAY: Duration = Duration::from_millis(10); } +const PACK_PROGRESS_TYPE: &str = "pack-progress"; + +/// A progress report emitted by a warm-up command on stdout. +/// +/// The command must output JSON Lines. For now only `type: "pack-progress"` is supported. +#[derive(Debug, Deserialize)] +struct PackProgress { + /// The message type. Must be `"pack-progress"`. + #[serde(rename = "type")] + ty: String, + + /// The number of packs the command reports as warm within this invocation. + warm: u64, +} + +/// Read JSON Lines progress reports from the warm-up command's stdout. +/// +/// Increments the shared progress bar by the delta between the last reported value and the new +/// one. Reports are ignored if they are not of type `pack-progress` or if `warm` is not larger than +/// the current value (progress is strictly increasing). +/// +/// The `invocation_size` is used to clamp the reported value; a command must never report more +/// packs than it received. +/// +/// Returns the final `warm` value reported by the command, or `0` if it emitted no protocol lines. +fn read_progress_output(reader: R, invocation_size: u64, progress: &Progress) -> u64 { + let reader = BufReader::new(reader); + let mut current: u64 = 0; + + for line in reader.lines().map_while(Result::ok) { + let report: PackProgress = match serde_json::from_str(&line) { + Ok(report) => report, + Err(_) => { + // For debugging/auditing purposes, log non-JSON stdout lines at info level. + // Empty lines are ignored. + if !line.is_empty() { + info!("[warmup] {line}"); + } + continue; + } + }; + + if report.ty != PACK_PROGRESS_TYPE { + continue; + } + + let warm = report.warm.min(invocation_size); + if warm > current { + progress.inc(warm - current); + current = warm; + } + } + + current +} + +/// On successful completion of a warm-up invocation, advance the progress bar by the packs that +/// were not already reported via the protocol. +fn finalize_progress(current: u64, invocation_size: u64, progress: &Progress) { + if current < invocation_size { + progress.inc(invocation_size - current); + } +} + /// Configuration for retrying executing a command that the operating system reports is busy. /// We believe this is a transient race condition that happens during unit tests when a program /// is created and then immediately executed. @@ -212,31 +277,48 @@ fn warm_up_batch_singular( debug!("spawning {command:?} for id {id:?}..."); - let child = (|| Command::new(command.command()).args(&args).spawn()) - .retry(execute_cmd_retry()) - .when(|err| err.kind() == io::ErrorKind::ExecutableFileBusy) - .notify(|err, duration| { - debug!("spawn failed with ETXTBSY, retrying in {duration:?}: {err}"); - }) - .call() - .map_err(|err| { - RusticError::with_source( - ErrorKind::ExternalCommand, - "Error in spawning warm-up command `{command}`.", - err, - ) - .attach_context("command", command.to_string()) - .attach_context("id", &id) - .attach_context("type", format!("{ty:?}")) - })?; + let child = (|| { + Command::new(command.command()) + .args(&args) + .stdout(Stdio::piped()) + .spawn() + }) + .retry(execute_cmd_retry()) + .when(|err| err.kind() == io::ErrorKind::ExecutableFileBusy) + .notify(|err, duration| { + debug!("spawn failed with ETXTBSY, retrying in {duration:?}: {err}"); + }) + .call() + .map_err(|err| { + RusticError::with_source( + ErrorKind::ExternalCommand, + "Error in spawning warm-up command `{command}`.", + err, + ) + .attach_context("command", command.to_string()) + .attach_context("id", &id) + .attach_context("type", format!("{ty:?}")) + })?; Ok((child, id)) }) .collect::>>()?; + // Spawn a reader for each child's stdout while the commands run concurrently. + let mut readers = Vec::with_capacity(children.len()); + for (mut child, id) in children { + let stdout = child.stdout.take().expect("stdout was piped"); + let progress = progress.clone(); + let handle = thread::spawn(move || { + // Singular mode handles one pack per command instance. + read_progress_output(stdout, 1, &progress) + }); + readers.push((child, id, handle)); + } + let mut failed_ids = Vec::new(); - for (mut child, id) in children { + for (mut child, id, handle) in readers { debug!("waiting for warm-up command for id {id}..."); let status = child.wait().map_err(|err| { @@ -250,11 +332,13 @@ fn warm_up_batch_singular( .attach_context("type", format!("{ty:?}")) })?; + let current = handle.join().unwrap(); + if !status.success() { failed_ids.push((id, status)); + } else { + finalize_progress(current, 1, progress); } - - progress.inc(1); } if !failed_ids.is_empty() { @@ -325,22 +409,46 @@ fn warm_up_batch_plural( debug!("calling {command:?} with {} id(s)...", batch.len()); - let status = (|| Command::new(command.command()).args(&args).status()) - .retry(execute_cmd_retry()) - .when(|err| err.kind() == io::ErrorKind::ExecutableFileBusy) - .notify(|err, duration| { - debug!("spawn failed with ETXTBSY, retrying in {duration:?}: {err}"); - }) - .call() - .map_err(|err| { - RusticError::with_source( - ErrorKind::ExternalCommand, - "Error in executing warm-up command `{command}`.", - err, - ) - .attach_context("command", command.to_string()) - .attach_context("type", format!("{ty:?}")) - })?; + let invocation_size = batch.len() as u64; + + let mut child = (|| { + Command::new(command.command()) + .args(&args) + .stdout(Stdio::piped()) + .spawn() + }) + .retry(execute_cmd_retry()) + .when(|err| err.kind() == io::ErrorKind::ExecutableFileBusy) + .notify(|err, duration| { + debug!("spawn failed with ETXTBSY, retrying in {duration:?}: {err}"); + }) + .call() + .map_err(|err| { + RusticError::with_source( + ErrorKind::ExternalCommand, + "Error in executing warm-up command `{command}`.", + err, + ) + .attach_context("command", command.to_string()) + .attach_context("type", format!("{ty:?}")) + })?; + + let stdout = child.stdout.take().expect("stdout was piped"); + let progress_for_thread = progress.clone(); + let handle = + thread::spawn(move || read_progress_output(stdout, invocation_size, &progress_for_thread)); + + let status = child.wait().map_err(|err| { + RusticError::with_source( + ErrorKind::ExternalCommand, + "Error waiting for warm-up command `{command}`.", + err, + ) + .attach_context("command", command.to_string()) + .attach_context("type", format!("{ty:?}")) + })?; + + let current = handle.join().unwrap(); if !status.success() { return Err(RusticError::new( @@ -356,7 +464,8 @@ fn warm_up_batch_plural( .attach_context("type", format!("{ty:?}"))); } - progress.inc(batch.len() as u64); + finalize_progress(current, invocation_size, progress); + Ok(()) } diff --git a/crates/core/tests/warm_up.rs b/crates/core/tests/warm_up.rs index cc9ca1e7..df7b9784 100644 --- a/crates/core/tests/warm_up.rs +++ b/crates/core/tests/warm_up.rs @@ -5,6 +5,7 @@ use std::{ io::Read, path::Path, sync::Arc, + sync::atomic::{AtomicU64, Ordering}, }; use anyhow::Result; @@ -12,7 +13,10 @@ use rstest::rstest; use serde::{Deserialize, Serialize}; use tempfile::tempdir; -use rustic_core::{CommandInput, Id, RepositoryBackends, RepositoryOptions, repofile::PackId}; +use rustic_core::{ + CommandInput, Id, Progress, ProgressBars, ProgressType, RepositoryBackends, RepositoryOptions, + RusticProgress, repofile::PackId, +}; use rustic_testing::backend::in_memory_backend::InMemoryBackend; // Test constants @@ -59,31 +63,7 @@ done "#, ); - // Write the script and sync to disk to avoid "Text file busy" errors - { - use std::io::Write; - let mut file = File::create(&script_path)?; - file.write_all(script_content.as_bytes())?; - file.sync_all()?; - // Explicitly drop the file handle before setting permissions - drop(file); - } - - // Make script executable - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let mut perms = fs::metadata(&script_path)?.permissions(); - perms.set_mode(0o755); - fs::set_permissions(&script_path, perms)?; - - // Sync the parent directory to ensure metadata changes are committed - if let Some(parent) = script_path.parent() - && let Ok(dir_file) = File::open(parent) - { - let _ = dir_file.sync_all(); - } - } + write_executable_script(&script_path, &script_content)?; let command: CommandInput = script_path.to_string_lossy().to_string().parse()?; @@ -102,6 +82,33 @@ fn create_failing_script(log_dir: &Path) -> Result<(tempfile::TempDir, CommandIn create_test_script_with_exit_code(log_dir, 1) } +/// Write a script to disk and mark it executable. +#[cfg(not(windows))] +fn write_executable_script(path: &Path, content: &str) -> Result<()> { + { + use std::io::Write; + let mut file = File::create(path)?; + file.write_all(content.as_bytes())?; + file.sync_all()?; + } + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mut perms = fs::metadata(path)?.permissions(); + perms.set_mode(0o755); + fs::set_permissions(path, perms)?; + + if let Some(parent) = path.parent() + && let Ok(dir_file) = File::open(parent) + { + let _ = dir_file.sync_all(); + } + } + + Ok(()) +} + /// Helper to parse log files in a directory and extract call count and arguments #[cfg(not(windows))] fn parse_log_files(log_dir: &Path) -> Result<(usize, Vec>)> { @@ -714,3 +721,171 @@ fn test_warm_up_backend_path_mode() -> Result<()> { Ok(()) } + +/// Progress implementation that records total increments. +#[derive(Debug, Clone)] +struct RecordingProgressBars { + pos: Arc, +} + +impl ProgressBars for RecordingProgressBars { + fn progress(&self, _progress_type: ProgressType, _prefix: &str) -> Progress { + Progress::new(RecordingProgress { + pos: self.pos.clone(), + }) + } +} + +#[derive(Debug, Clone)] +struct RecordingProgress { + pos: Arc, +} + +impl RusticProgress for RecordingProgress { + fn is_hidden(&self) -> bool { + false + } + fn set_length(&self, _len: u64) {} + fn set_title(&self, _title: &str) {} + fn inc(&self, inc: u64) { + let _ = self.pos.fetch_add(inc, Ordering::Relaxed); + } + fn finish(&self) {} +} + +#[cfg(not(windows))] +#[test] +fn test_warm_up_progress_protocol_plural() -> Result<()> { + let dir = tempdir()?; + let script_path = dir.path().join("progress.sh"); + + let script_content = r#"#!/usr/bin/env bash +warm=0 +for arg in "$@"; do + warm=$((warm+1)) + echo "{\"type\":\"pack-progress\",\"warm\":$warm}" +done +"#; + write_executable_script(&script_path, script_content)?; + + let cmd_str = format!("{} %ids", script_path.to_string_lossy()); + let command: CommandInput = cmd_str.parse()?; + + let pos = Arc::new(AtomicU64::new(0)); + let repo = rustic_core::Repository::new_with_progress( + &RepositoryOptions::default() + .warm_up_command(command) + .warm_up_batch(5usize), + &RepositoryBackends::new(Arc::new(InMemoryBackend::new()), None), + RecordingProgressBars { pos: pos.clone() }, + )?; + + let num_packs = 12; + let pack_ids = create_test_ids(num_packs); + + repo.warm_up(pack_ids.iter().copied())?; + + assert_eq!(pos.load(Ordering::Relaxed), num_packs as u64); + + Ok(()) +} + +#[cfg(not(windows))] +#[test] +fn test_warm_up_progress_protocol_singular() -> Result<()> { + let dir = tempdir()?; + let script_path = dir.path().join("progress.sh"); + + let script_content = r#"#!/usr/bin/env bash +echo '{"type":"pack-progress","warm":1}' +"#; + write_executable_script(&script_path, script_content)?; + + let cmd_str = format!("{} %id", script_path.to_string_lossy()); + let command: CommandInput = cmd_str.parse()?; + + let pos = Arc::new(AtomicU64::new(0)); + let repo = rustic_core::Repository::new_with_progress( + &RepositoryOptions::default() + .warm_up_command(command) + .warm_up_batch(1usize), + &RepositoryBackends::new(Arc::new(InMemoryBackend::new()), None), + RecordingProgressBars { pos: pos.clone() }, + )?; + + let num_packs = 5; + let pack_ids = create_test_ids(num_packs); + + repo.warm_up(pack_ids.iter().copied())?; + + assert_eq!(pos.load(Ordering::Relaxed), num_packs as u64); + + Ok(()) +} + +#[cfg(not(windows))] +#[test] +fn test_warm_up_progress_protocol_fallback() -> Result<()> { + let dir = tempdir()?; + let script_path = dir.path().join("silent.sh"); + + let script_content = "#!/usr/bin/env bash\n"; + write_executable_script(&script_path, script_content)?; + + let cmd_str = format!("{} %ids", script_path.to_string_lossy()); + let command: CommandInput = cmd_str.parse()?; + + let pos = Arc::new(AtomicU64::new(0)); + let repo = rustic_core::Repository::new_with_progress( + &RepositoryOptions::default() + .warm_up_command(command) + .warm_up_batch(4usize), + &RepositoryBackends::new(Arc::new(InMemoryBackend::new()), None), + RecordingProgressBars { pos: pos.clone() }, + )?; + + let num_packs = 9; + let pack_ids = create_test_ids(num_packs); + + repo.warm_up(pack_ids.iter().copied())?; + + assert_eq!(pos.load(Ordering::Relaxed), num_packs as u64); + + Ok(()) +} + +#[cfg(not(windows))] +#[test] +fn test_warm_up_progress_protocol_ignores_invalid_lines() -> Result<()> { + let dir = tempdir()?; + let script_path = dir.path().join("noisy.sh"); + + let script_content = r#"#!/usr/bin/env bash +echo 'this is not json' +echo '{}' +echo '{"type":"other","warm":100}' +echo '{"type":"pack-progress","warm":999}' +"#; + write_executable_script(&script_path, script_content)?; + + let cmd_str = format!("{} %ids", script_path.to_string_lossy()); + let command: CommandInput = cmd_str.parse()?; + + let pos = Arc::new(AtomicU64::new(0)); + let repo = rustic_core::Repository::new_with_progress( + &RepositoryOptions::default() + .warm_up_command(command) + .warm_up_batch(3usize), + &RepositoryBackends::new(Arc::new(InMemoryBackend::new()), None), + RecordingProgressBars { pos: pos.clone() }, + )?; + + let num_packs = 3; + let pack_ids = create_test_ids(num_packs); + + repo.warm_up(pack_ids.iter().copied())?; + + assert_eq!(pos.load(Ordering::Relaxed), num_packs as u64); + + Ok(()) +}