pub mod replay;
pub mod writer;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use anyhow::{Context, Result, anyhow, bail};
use serde_json::json;
use crate::cli::{RecordArgs, RecordCompression, RecordSource, SnapshotIncludes};
use crate::common::config_file::RecordSettings;
use crate::snapshot::serializers::write_frame_json;
use crate::snapshot::{
DefaultSnapshotCollector, SNAPSHOT_SCHEMA_VERSION, Snapshot, SnapshotCollector,
SnapshotOptions, collect_once,
};
use writer::{Codec, RotatingWriter};
const INDEX_EVERY_N_FRAMES: u64 = 1000;
pub const DEFAULT_BASENAME: &str = "all-smi-record.ndjson.zst";
pub struct RecorderOptions {
pub output: PathBuf,
pub interval: Duration,
pub duration: Option<Duration>,
pub source: RecordSource,
pub hosts: Vec<String>,
pub hostfile: Option<String>,
pub includes: SnapshotIncludes,
pub max_size: u64,
pub max_files: u32,
pub codec: Codec,
}
impl RecorderOptions {
#[allow(dead_code)]
pub fn from_args(args: &RecordArgs) -> Result<Self> {
Self::from_args_with_settings(args, None)
}
pub fn from_args_with_settings(
args: &RecordArgs,
settings: Option<&RecordSettings>,
) -> Result<Self> {
let includes = args
.includes()
.map_err(|msg| anyhow!("invalid --include: {msg}"))?;
if includes.is_empty() {
bail!("at least one section must be requested via --include");
}
let duration = parse_duration(&args.duration).context("invalid --duration")?;
let max_size = parse_byte_size(&args.max_size).context("invalid --max-size")?;
if args.max_files == 0 {
bail!("--max-files must be >= 1");
}
let hosts: Vec<String> = args.hosts.clone().unwrap_or_default();
if args.source == RecordSource::Remote && hosts.is_empty() && args.hostfile.is_none() {
bail!("--source=remote requires --hosts or --hostfile");
}
let compress_override: Option<RecordCompression> = args.compress.or_else(|| {
settings.map(|s| s.compress.as_str()).and_then(|c| match c {
"zstd" => Some(RecordCompression::Zstd),
"gzip" => Some(RecordCompression::Gzip),
"none" => Some(RecordCompression::None),
_ => None, })
});
let output = match &args.output {
Some(p) => p.clone(),
None => {
let configured = settings
.and_then(|s| s.output_dir.as_deref())
.map(str::trim)
.filter(|s| !s.is_empty());
let base = match configured {
Some(d) => Some(crate::common::paths::expand_tilde(std::path::Path::new(d))),
None => crate::common::paths::cache_dir().map(|d| d.join("records")),
};
match base {
Some(b) => b.join(DEFAULT_BASENAME),
None => std::path::PathBuf::from(DEFAULT_BASENAME),
}
}
};
let codec = Codec::detect(&output, compress_override);
if compress_override.is_some()
&& let Some(warn) = codec_extension_mismatch(&output, compress_override)
{
eprintln!("warning: {warn}");
}
Ok(Self {
output,
interval: Duration::from_secs(args.interval.max(1)),
duration,
source: args.source,
hosts,
hostfile: args.hostfile.clone(),
includes,
max_size,
max_files: args.max_files,
codec,
})
}
}
pub async fn run(opts: RecorderOptions) -> Result<()> {
let stop = Arc::new(AtomicBool::new(false));
install_signal_handlers(stop.clone());
match opts.source {
RecordSource::Local => run_local(opts, stop).await,
RecordSource::Remote => run_remote(opts, stop).await,
}
}
async fn run_local(opts: RecorderOptions, stop: Arc<AtomicBool>) -> Result<()> {
let collector = Arc::new(DefaultSnapshotCollector::new());
let hosts = vec![collector.hostname()];
let mut writer =
RotatingWriter::new(&opts.output, opts.codec, opts.max_size, opts.max_files)
.with_context(|| format!("failed to open recording file {}", opts.output.display()))?;
write_header(&mut writer, &opts, &hosts)?;
record_loop(&mut writer, &opts, stop, move |includes, timeout| {
let c = collector.clone();
let inc = *includes;
async move { collect_once(c, &inc, timeout).await }
})
.await?;
writer.finish().context("failed to finalize recording")?;
Ok(())
}
async fn run_remote(opts: RecorderOptions, stop: Arc<AtomicBool>) -> Result<()> {
use crate::view::data_collection::{
CollectionConfig, DataCollectionStrategy, RemoteCollectorBuilder,
};
let mut builder = RemoteCollectorBuilder::new().with_hosts(opts.hosts.clone());
if let Some(file) = opts.hostfile.as_deref() {
builder = builder
.load_hosts_from_file(file)
.with_context(|| format!("failed to load hostfile {file}"))?;
}
let collector = builder.build();
let mut writer =
RotatingWriter::new(&opts.output, opts.codec, opts.max_size, opts.max_files)
.with_context(|| format!("failed to open recording file {}", opts.output.display()))?;
write_header(&mut writer, &opts, &opts.hosts)?;
let start = std::time::Instant::now();
let mut seq: u64 = 0;
while !stop.load(Ordering::Relaxed) {
if let Some(limit) = opts.duration
&& start.elapsed() >= limit
{
break;
}
let config = CollectionConfig {
interval: opts.interval.as_secs(),
first_iteration: false,
hosts: opts.hosts.clone(),
};
match collector.collect(&config).await {
Ok(data) => {
let snap = snapshot_from_collection_data(&data, &opts);
write_data_frame(&mut writer, &snap, seq)?;
seq += 1;
}
Err(e) => {
tracing::warn!(error = %e, "record: remote scrape failed");
}
}
sleep_until_next(&opts, start, seq, &stop).await;
}
writer.finish().context("failed to finalize recording")?;
Ok(())
}
async fn record_loop<F, Fut>(
writer: &mut RotatingWriter,
opts: &RecorderOptions,
stop: Arc<AtomicBool>,
mut collect: F,
) -> Result<()>
where
F: FnMut(&SnapshotIncludes, Duration) -> Fut,
Fut: std::future::Future<Output = Snapshot>,
{
let start = std::time::Instant::now();
let mut seq: u64 = 0;
let reader_timeout = Duration::from_millis(5_000);
while !stop.load(Ordering::Relaxed) {
if let Some(limit) = opts.duration
&& start.elapsed() >= limit
{
break;
}
let snap = collect(&opts.includes, reader_timeout).await;
write_data_frame(writer, &snap, seq)?;
seq += 1;
sleep_until_next(opts, start, seq, &stop).await;
}
Ok(())
}
async fn sleep_until_next(
opts: &RecorderOptions,
start: std::time::Instant,
seq: u64,
stop: &Arc<AtomicBool>,
) {
let target = start + opts.interval.saturating_mul(seq as u32);
let now = std::time::Instant::now();
if target > now {
let remaining = target.duration_since(now);
let mut slept = Duration::ZERO;
let step = Duration::from_millis(100);
while slept < remaining && !stop.load(Ordering::Relaxed) {
let chunk = step.min(remaining - slept);
tokio::time::sleep(chunk).await;
slept += chunk;
}
}
}
fn write_header(
writer: &mut RotatingWriter,
opts: &RecorderOptions,
hosts: &[String],
) -> Result<()> {
let header = json!({
"schema": SNAPSHOT_SCHEMA_VERSION,
"header": true,
"interval_ms": opts.interval.as_millis() as u64,
"hosts": hosts,
"all_smi_version": env!("CARGO_PKG_VERSION"),
});
let mut line = serde_json::to_string(&header).context("failed to serialize header frame")?;
line.push('\n');
writer.write_line(line.as_bytes())?;
Ok(())
}
fn write_data_frame(writer: &mut RotatingWriter, snapshot: &Snapshot, seq: u64) -> Result<()> {
let mut buf: Vec<u8> = Vec::with_capacity(4096);
write_frame_json(&mut buf, snapshot).context("failed to serialize data frame")?;
writer.write_line(&buf)?;
if seq > 0 && seq.is_multiple_of(INDEX_EVERY_N_FRAMES) {
let idx = json!({
"schema": SNAPSHOT_SCHEMA_VERSION,
"index": true,
"seq": seq,
"byte_offset": writer.active_bytes(),
});
let mut line = serde_json::to_string(&idx).context("failed to serialize index frame")?;
line.push('\n');
writer.write_line(line.as_bytes())?;
}
Ok(())
}
fn snapshot_from_collection_data(
data: &crate::view::data_collection::strategy::CollectionData,
opts: &RecorderOptions,
) -> Snapshot {
let mut snap = Snapshot {
schema: SNAPSHOT_SCHEMA_VERSION,
timestamp: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
hostname: "remote".to_string(),
gpus: None,
cpus: None,
memory: None,
chassis: None,
processes: None,
storage: None,
errors: Vec::new(),
};
if opts.includes.gpu {
snap.gpus = Some(data.gpu_info.clone());
}
if opts.includes.cpu {
snap.cpus = Some(data.cpu_info.clone());
}
if opts.includes.memory {
snap.memory = Some(data.memory_info.clone());
}
if opts.includes.chassis {
snap.chassis = Some(data.chassis_info.clone());
}
if opts.includes.process {
snap.processes = Some(data.process_info.clone());
}
snap
}
fn install_signal_handlers(stop: Arc<AtomicBool>) {
{
let stop = stop.clone();
tokio::spawn(async move {
let _ = tokio::signal::ctrl_c().await;
stop.store(true, Ordering::Relaxed);
});
}
#[cfg(unix)]
{
let stop = stop.clone();
tokio::spawn(async move {
if let Ok(mut sig) =
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
{
sig.recv().await;
stop.store(true, Ordering::Relaxed);
}
});
}
}
fn parse_duration(s: &str) -> Result<Option<Duration>> {
let trimmed = s.trim();
if trimmed.is_empty() {
return Ok(None);
}
if trimmed == "0" {
return Ok(None);
}
let (num_part, unit_secs): (&str, u64) = if let Some(stripped) = trimmed.strip_suffix('s') {
(stripped, 1)
} else if let Some(stripped) = trimmed.strip_suffix('m') {
(stripped, 60)
} else if let Some(stripped) = trimmed.strip_suffix('h') {
(stripped, 3_600)
} else if let Some(stripped) = trimmed.strip_suffix('d') {
(stripped, 86_400)
} else {
(trimmed, 1)
};
let n: u64 = num_part
.parse()
.map_err(|_| anyhow!("expected integer, got `{trimmed}`"))?;
Ok(Some(Duration::from_secs(n.saturating_mul(unit_secs))))
}
fn parse_byte_size(s: &str) -> Result<u64> {
let trimmed = s.trim();
if trimmed.is_empty() {
return Ok(0);
}
let (num_part, mul): (&str, u64) = if let Some(stripped) = trimmed.strip_suffix('K') {
(stripped, 1 << 10)
} else if let Some(stripped) = trimmed.strip_suffix('M') {
(stripped, 1 << 20)
} else if let Some(stripped) = trimmed.strip_suffix('G') {
(stripped, 1 << 30)
} else {
(trimmed, 1)
};
let n: u64 = num_part
.parse()
.map_err(|_| anyhow!("expected integer, got `{trimmed}`"))?;
Ok(n.saturating_mul(mul))
}
fn codec_extension_mismatch(path: &Path, compress: Option<RecordCompression>) -> Option<String> {
let ext = path
.extension()
.and_then(|e| e.to_str())
.map(|e| e.to_ascii_lowercase());
let ext_codec = match ext.as_deref() {
Some("zst") => Some("zstd"),
Some("gz") => Some("gzip"),
_ => Some("plain"),
};
let forced = match compress? {
RecordCompression::Zstd => "zstd",
RecordCompression::Gzip => "gzip",
RecordCompression::None => "plain",
};
match (ext_codec, forced) {
(Some(e), f) if e != f => Some(format!(
"--compress={f} overrides file extension `.{}` which suggests {e}",
ext.unwrap_or_default()
)),
_ => None,
}
}
#[allow(dead_code)]
fn _type_check_snapshot_reexport(_: SnapshotOptions) {}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_duration_accepts_suffixes() {
assert_eq!(parse_duration("0").unwrap(), None);
assert_eq!(parse_duration("30").unwrap(), Some(Duration::from_secs(30)));
assert_eq!(
parse_duration("30s").unwrap(),
Some(Duration::from_secs(30))
);
assert_eq!(
parse_duration("5m").unwrap(),
Some(Duration::from_secs(300))
);
assert_eq!(
parse_duration("1h").unwrap(),
Some(Duration::from_secs(3_600))
);
assert_eq!(
parse_duration("2d").unwrap(),
Some(Duration::from_secs(172_800))
);
}
#[test]
fn parse_duration_rejects_junk() {
assert!(parse_duration("abc").is_err());
assert!(parse_duration("5x").is_err());
}
#[test]
fn parse_byte_size_accepts_suffixes() {
assert_eq!(parse_byte_size("0").unwrap(), 0);
assert_eq!(parse_byte_size("1024").unwrap(), 1024);
assert_eq!(parse_byte_size("1K").unwrap(), 1024);
assert_eq!(parse_byte_size("1M").unwrap(), 1 << 20);
assert_eq!(parse_byte_size("2G").unwrap(), 2 << 30);
}
fn record_args(output: Option<PathBuf>) -> RecordArgs {
RecordArgs {
output,
interval: 3,
duration: "0".to_string(),
source: RecordSource::Local,
hosts: None,
hostfile: None,
include: vec![
"gpu".to_string(),
"cpu".to_string(),
"memory".to_string(),
"chassis".to_string(),
],
max_size: "100M".to_string(),
max_files: 10,
compress: None,
}
}
#[test]
fn resolve_output_no_cli_no_config() {
let args = record_args(None);
let opts = RecorderOptions::from_args_with_settings(&args, None).unwrap();
let expected = match crate::common::paths::cache_dir() {
Some(c) => c.join("records").join(DEFAULT_BASENAME),
None => PathBuf::from(DEFAULT_BASENAME),
};
assert_eq!(opts.output, expected);
}
#[test]
fn resolve_output_no_cli_no_config_uses_platform_cache_dir() {
let Some(cache) = crate::common::paths::cache_dir() else {
return;
};
let args = record_args(None);
let opts = RecorderOptions::from_args_with_settings(&args, None).unwrap();
assert_eq!(opts.output, cache.join("records").join(DEFAULT_BASENAME));
}
#[test]
fn resolve_output_no_cli_with_compiled_record_settings_uses_cache_dir() {
let args = record_args(None);
let settings = crate::common::config_file::Settings::default().record;
let opts = RecorderOptions::from_args_with_settings(&args, Some(&settings)).unwrap();
let expected = match crate::common::paths::cache_dir() {
Some(c) => c.join("records").join(DEFAULT_BASENAME),
None => PathBuf::from(DEFAULT_BASENAME),
};
assert_eq!(opts.output, expected);
}
#[test]
fn resolve_output_no_cli_with_config_dir() {
let args = record_args(None);
let settings = RecordSettings {
output_dir: Some("/tmp/all-smi-records".to_string()),
compress: "zstd".to_string(),
};
let opts = RecorderOptions::from_args_with_settings(&args, Some(&settings)).unwrap();
assert_eq!(
opts.output,
PathBuf::from("/tmp/all-smi-records/all-smi-record.ndjson.zst")
);
}
#[test]
fn resolve_output_explicit_matching_basename_is_honored() {
let args = record_args(Some(PathBuf::from(DEFAULT_BASENAME)));
let settings = RecordSettings {
output_dir: Some("/tmp/all-smi-records".to_string()),
compress: "zstd".to_string(),
};
let opts = RecorderOptions::from_args_with_settings(&args, Some(&settings)).unwrap();
assert_eq!(opts.output, PathBuf::from(DEFAULT_BASENAME));
}
#[test]
fn resolve_output_explicit_absolute_path_is_honored() {
let explicit = PathBuf::from("/var/log/cluster.ndjson.zst");
let args = record_args(Some(explicit.clone()));
let opts = RecorderOptions::from_args_with_settings(&args, None).unwrap();
assert_eq!(opts.output, explicit);
let args = record_args(Some(explicit.clone()));
let settings = RecordSettings {
output_dir: Some("/tmp/all-smi-records".to_string()),
compress: "zstd".to_string(),
};
let opts = RecorderOptions::from_args_with_settings(&args, Some(&settings)).unwrap();
assert_eq!(opts.output, explicit);
}
#[test]
fn resolve_output_whitespace_only_config_dir_falls_through() {
let args = record_args(None);
let settings = RecordSettings {
output_dir: Some(" ".to_string()),
compress: "zstd".to_string(),
};
let opts = RecorderOptions::from_args_with_settings(&args, Some(&settings)).unwrap();
let s = opts.output.to_string_lossy();
assert!(
!s.starts_with(" "),
"whitespace-only config_dir leaked into output: {s}"
);
}
#[test]
fn resolve_output_no_cli_with_tilde_config_dir() {
let args = record_args(None);
let settings = RecordSettings {
output_dir: Some("~/my-records".to_string()),
compress: "zstd".to_string(),
};
let opts = RecorderOptions::from_args_with_settings(&args, Some(&settings)).unwrap();
let expected = crate::common::paths::expand_tilde(std::path::Path::new("~/my-records"))
.join(DEFAULT_BASENAME);
assert_eq!(opts.output, expected);
}
}