304 lines
10 KiB
Rust
304 lines
10 KiB
Rust
|
//! Tarball assembly for a finalized run.
|
||
|
|
//!
|
||
|
|
//! Walks `{root}/{run_id}/` and packs it into
|
||
|
|
//! `{root}/bundles/{run_id}.tar.gz`. Per-node directories are
|
||
|
|
//! relabeled from `{node_id}` to a friendlier `{role}` (or
|
||
|
|
//! `{role}-{stage_index}`) when the boot record gave us a role. The
|
||
|
|
//! mapping is recorded in `MANIFEST.json` so the full hex is always
|
||
|
|
//! recoverable.
|
||
|
|
|
||
|
|
use std::collections::BTreeMap;
|
||
|
|
use std::fs::File;
|
||
|
|
use std::io;
|
||
|
|
use std::path::{Path, PathBuf};
|
||
|
|
|
||
|
|
use flate2::Compression;
|
||
|
|
use flate2::write::GzEncoder;
|
||
|
|
use serde_json::Value;
|
||
|
|
|
||
|
|
use super::protocol::{Manifest, ManifestNode};
|
||
|
|
use super::state::CollectorState;
|
||
|
|
|
||
|
|
pub fn assemble(state: &CollectorState, run_id: &str) -> io::Result<PathBuf> {
|
||
|
|
let run_dir = state.run_dir(run_id);
|
||
|
|
if !run_dir.is_dir() {
|
||
|
|
return Err(io::Error::new(
|
||
|
|
io::ErrorKind::NotFound,
|
||
|
|
format!("no records on disk for run_id {run_id}"),
|
||
|
|
));
|
||
|
|
}
|
||
|
|
let bundles_dir = state.bundles_dir();
|
||
|
|
std::fs::create_dir_all(&bundles_dir)?;
|
||
|
|
let bundle_path = state.bundle_path(run_id);
|
||
|
|
|
||
|
|
let stats = state.run_stats(run_id);
|
||
|
|
let labels = build_labels(&stats);
|
||
|
|
let manifest = build_manifest(run_id, &stats, &labels);
|
||
|
|
|
||
|
|
let file = File::create(&bundle_path)?;
|
||
|
|
let gz = GzEncoder::new(file, Compression::default());
|
||
|
|
let mut tar = tar::Builder::new(gz);
|
||
|
|
tar.mode(tar::HeaderMode::Deterministic);
|
||
|
|
|
||
|
|
// Root entry for the run.
|
||
|
|
append_dir(&mut tar, run_id)?;
|
||
|
|
|
||
|
|
// MANIFEST.json at tar root for this run.
|
||
|
|
let manifest_bytes = serde_json::to_vec_pretty(&manifest)
|
||
|
|
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
|
||
|
|
append_bytes(
|
||
|
|
&mut tar,
|
||
|
|
&format!("{run_id}/MANIFEST.json"),
|
||
|
|
&manifest_bytes,
|
||
|
|
)?;
|
||
|
|
|
||
|
|
// Walk per-node directories in label order so the tarball is
|
||
|
|
// deterministic for tests / cross-run diffs.
|
||
|
|
let mut ordered: Vec<(&String, &String)> = labels.iter().collect();
|
||
|
|
ordered.sort_by(|a, b| a.1.cmp(b.1).then(a.0.cmp(b.0)));
|
||
|
|
for (node_id, label) in ordered {
|
||
|
|
let node_src = state.node_dir(run_id, node_id);
|
||
|
|
if !node_src.is_dir() {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let tar_node = format!("{run_id}/{label}");
|
||
|
|
append_dir(&mut tar, &tar_node)?;
|
||
|
|
append_node_dir(&mut tar, &node_src, &tar_node)?;
|
||
|
|
}
|
||
|
|
|
||
|
|
tar.finish()?;
|
||
|
|
Ok(bundle_path)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn append_node_dir(
|
||
|
|
tar: &mut tar::Builder<GzEncoder<File>>,
|
||
|
|
src: &Path,
|
||
|
|
dst_prefix: &str,
|
||
|
|
) -> io::Result<()> {
|
||
|
|
// We organize files by kind into kind-named subdirs to match
|
||
|
|
// the bundle layout in DIAGNOSTICS_PLAN.md (snapshots/, events/).
|
||
|
|
// Boot and finalize records become single files at the node root.
|
||
|
|
let mut boots: Vec<PathBuf> = Vec::new();
|
||
|
|
let mut events: Vec<PathBuf> = Vec::new();
|
||
|
|
let mut snapshots: Vec<PathBuf> = Vec::new();
|
||
|
|
let mut finalizes: Vec<PathBuf> = Vec::new();
|
||
|
|
for entry in std::fs::read_dir(src)? {
|
||
|
|
let entry = entry?;
|
||
|
|
let path = entry.path();
|
||
|
|
if !path.is_file() {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let name = match path.file_name().and_then(|n| n.to_str()) {
|
||
|
|
Some(n) => n,
|
||
|
|
None => continue,
|
||
|
|
};
|
||
|
|
if name.starts_with("boot-") {
|
||
|
|
boots.push(path);
|
||
|
|
} else if name.starts_with("events-") {
|
||
|
|
events.push(path);
|
||
|
|
} else if name.starts_with("snapshot-") {
|
||
|
|
snapshots.push(path);
|
||
|
|
} else if name.starts_with("finalize-") {
|
||
|
|
finalizes.push(path);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
boots.sort();
|
||
|
|
events.sort();
|
||
|
|
snapshots.sort();
|
||
|
|
finalizes.sort();
|
||
|
|
|
||
|
|
// Promote the latest boot/finalize to a single file at the node
|
||
|
|
// root — the format in DIAGNOSTICS_PLAN.md expects exactly one of
|
||
|
|
// each. Earlier boot retries (if any) stay under `boot/` for
|
||
|
|
// forensic value.
|
||
|
|
if let Some(latest_boot) = boots.last() {
|
||
|
|
append_file(tar, latest_boot, &format!("{dst_prefix}/boot.json"))?;
|
||
|
|
}
|
||
|
|
if let Some(latest_finalize) = finalizes.last() {
|
||
|
|
append_file(tar, latest_finalize, &format!("{dst_prefix}/finalize.json"))?;
|
||
|
|
}
|
||
|
|
if boots.len() > 1 {
|
||
|
|
append_dir(tar, &format!("{dst_prefix}/boot"))?;
|
||
|
|
for p in &boots[..boots.len() - 1] {
|
||
|
|
append_under(tar, p, &format!("{dst_prefix}/boot"))?;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if !events.is_empty() {
|
||
|
|
append_dir(tar, &format!("{dst_prefix}/events"))?;
|
||
|
|
for p in &events {
|
||
|
|
append_under(tar, p, &format!("{dst_prefix}/events"))?;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if !snapshots.is_empty() {
|
||
|
|
append_dir(tar, &format!("{dst_prefix}/snapshots"))?;
|
||
|
|
for p in &snapshots {
|
||
|
|
append_under(tar, p, &format!("{dst_prefix}/snapshots"))?;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if finalizes.len() > 1 {
|
||
|
|
append_dir(tar, &format!("{dst_prefix}/finalize"))?;
|
||
|
|
for p in &finalizes[..finalizes.len() - 1] {
|
||
|
|
append_under(tar, p, &format!("{dst_prefix}/finalize"))?;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn append_under(
|
||
|
|
tar: &mut tar::Builder<GzEncoder<File>>,
|
||
|
|
src: &Path,
|
||
|
|
dst_dir: &str,
|
||
|
|
) -> io::Result<()> {
|
||
|
|
let name = src
|
||
|
|
.file_name()
|
||
|
|
.and_then(|n| n.to_str())
|
||
|
|
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "non-utf8 filename"))?;
|
||
|
|
append_file(tar, src, &format!("{dst_dir}/{name}"))
|
||
|
|
}
|
||
|
|
|
||
|
|
fn append_file(
|
||
|
|
tar: &mut tar::Builder<GzEncoder<File>>,
|
||
|
|
src: &Path,
|
||
|
|
dst: &str,
|
||
|
|
) -> io::Result<()> {
|
||
|
|
let mut f = File::open(src)?;
|
||
|
|
let meta = f.metadata()?;
|
||
|
|
let mut header = tar::Header::new_gnu();
|
||
|
|
header.set_size(meta.len());
|
||
|
|
header.set_mode(0o644);
|
||
|
|
header.set_mtime(0);
|
||
|
|
header.set_entry_type(tar::EntryType::Regular);
|
||
|
|
header.set_cksum();
|
||
|
|
tar.append_data(&mut header, dst, &mut f)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn append_dir(tar: &mut tar::Builder<GzEncoder<File>>, dst: &str) -> io::Result<()> {
|
||
|
|
let mut header = tar::Header::new_gnu();
|
||
|
|
header.set_size(0);
|
||
|
|
header.set_mode(0o755);
|
||
|
|
header.set_mtime(0);
|
||
|
|
header.set_entry_type(tar::EntryType::Directory);
|
||
|
|
header.set_cksum();
|
||
|
|
let path = format!("{}/", dst.trim_end_matches('/'));
|
||
|
|
tar.append_data(&mut header, path, &mut io::empty())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn append_bytes(
|
||
|
|
tar: &mut tar::Builder<GzEncoder<File>>,
|
||
|
|
dst: &str,
|
||
|
|
bytes: &[u8],
|
||
|
|
) -> io::Result<()> {
|
||
|
|
let mut header = tar::Header::new_gnu();
|
||
|
|
header.set_size(bytes.len() as u64);
|
||
|
|
header.set_mode(0o644);
|
||
|
|
header.set_mtime(0);
|
||
|
|
header.set_entry_type(tar::EntryType::Regular);
|
||
|
|
header.set_cksum();
|
||
|
|
tar.append_data(&mut header, dst, bytes)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn build_labels(stats: &Option<super::state::RunStats>) -> BTreeMap<String, String> {
|
||
|
|
// Maps full hex node_id -> bundle-side directory label. Falls
|
||
|
|
// back to `node-{short}` when boot didn't supply enough.
|
||
|
|
let mut out = BTreeMap::new();
|
||
|
|
let Some(stats) = stats else { return out };
|
||
|
|
for (node_id, node) in &stats.nodes {
|
||
|
|
let label = label_for(node_id, node.identity.as_ref());
|
||
|
|
out.insert(node_id.clone(), label);
|
||
|
|
}
|
||
|
|
// Disambiguate label collisions (two nodes both calling themselves
|
||
|
|
// "orchestrator", for example) by suffixing with the short id.
|
||
|
|
let mut seen: BTreeMap<String, u32> = BTreeMap::new();
|
||
|
|
for label in out.values() {
|
||
|
|
*seen.entry(label.clone()).or_insert(0) += 1;
|
||
|
|
}
|
||
|
|
let collisions: Vec<String> = seen
|
||
|
|
.iter()
|
||
|
|
.filter(|(_, c)| **c > 1)
|
||
|
|
.map(|(k, _)| k.clone())
|
||
|
|
.collect();
|
||
|
|
for (node_id, label) in out.iter_mut() {
|
||
|
|
if collisions.contains(label) {
|
||
|
|
let short = node_id.chars().take(8).collect::<String>();
|
||
|
|
*label = format!("{label}-{short}");
|
||
|
|
}
|
||
|
|
}
|
||
|
|
out
|
||
|
|
}
|
||
|
|
|
||
|
|
fn label_for(node_id: &str, identity: Option<&Value>) -> String {
|
||
|
|
let short: String = node_id.chars().take(8).collect();
|
||
|
|
let Some(id) = identity else {
|
||
|
|
return format!("node-{short}");
|
||
|
|
};
|
||
|
|
let role = id.get("role").and_then(|r| match r {
|
||
|
|
Value::String(s) => Some(s.clone()),
|
||
|
|
Value::Object(o) => o.get("0").and_then(|v| v.as_str()).map(str::to_string),
|
||
|
|
Value::Array(arr) => arr.first().and_then(|v| v.as_str()).map(str::to_string),
|
||
|
|
_ => None,
|
||
|
|
});
|
||
|
|
let stage_index = id.get("stage_index").and_then(|v| v.as_u64());
|
||
|
|
match (role.as_deref(), stage_index) {
|
||
|
|
(Some("stage"), Some(idx)) => format!("stage-{idx}"),
|
||
|
|
(Some(r), _) if !r.is_empty() => r.to_string(),
|
||
|
|
_ => format!("node-{short}"),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn build_manifest(
|
||
|
|
run_id: &str,
|
||
|
|
stats: &Option<super::state::RunStats>,
|
||
|
|
labels: &BTreeMap<String, String>,
|
||
|
|
) -> Manifest {
|
||
|
|
let stats = match stats {
|
||
|
|
Some(s) => s,
|
||
|
|
None => {
|
||
|
|
return Manifest {
|
||
|
|
run_id: run_id.to_string(),
|
||
|
|
run_start_collector_ms: None,
|
||
|
|
run_end_collector_ms: None,
|
||
|
|
finalize_received: false,
|
||
|
|
nodes: Vec::new(),
|
||
|
|
};
|
||
|
|
}
|
||
|
|
};
|
||
|
|
let mut nodes: Vec<ManifestNode> = stats
|
||
|
|
.nodes
|
||
|
|
.iter()
|
||
|
|
.map(|(node_id, node)| {
|
||
|
|
let identity = node.identity.as_ref();
|
||
|
|
let role = identity.and_then(|id| id.get("role")).and_then(|r| match r {
|
||
|
|
Value::String(s) => Some(s.clone()),
|
||
|
|
Value::Object(o) => o.get("0").and_then(|v| v.as_str()).map(str::to_string),
|
||
|
|
Value::Array(arr) => arr.first().and_then(|v| v.as_str()).map(str::to_string),
|
||
|
|
_ => None,
|
||
|
|
});
|
||
|
|
let stage_index = identity
|
||
|
|
.and_then(|id| id.get("stage_index"))
|
||
|
|
.and_then(|v| v.as_u64())
|
||
|
|
.map(|v| v as u32);
|
||
|
|
ManifestNode {
|
||
|
|
node_id_hex: node_id.clone(),
|
||
|
|
label: labels
|
||
|
|
.get(node_id)
|
||
|
|
.cloned()
|
||
|
|
.unwrap_or_else(|| format!("node-{}", node_id.chars().take(8).collect::<String>())),
|
||
|
|
role,
|
||
|
|
stage_index,
|
||
|
|
boot_recorded: node.boot_recorded,
|
||
|
|
event_batches: node.event_batches,
|
||
|
|
snapshots: node.snapshots,
|
||
|
|
finalize_recorded: node.finalize_recorded,
|
||
|
|
}
|
||
|
|
})
|
||
|
|
.collect();
|
||
|
|
nodes.sort_by(|a, b| a.label.cmp(&b.label).then(a.node_id_hex.cmp(&b.node_id_hex)));
|
||
|
|
Manifest {
|
||
|
|
run_id: run_id.to_string(),
|
||
|
|
run_start_collector_ms: stats.run_start_collector_ms,
|
||
|
|
run_end_collector_ms: stats.run_end_collector_ms,
|
||
|
|
finalize_received: stats.finalize_received,
|
||
|
|
nodes,
|
||
|
|
}
|
||
|
|
}
|