From 603d113cd23a60dbd0f31b95716a044c24ada8d0 Mon Sep 17 00:00:00 2001 From: peytonr18 Date: Thu, 9 Jul 2026 16:29:43 -0700 Subject: [PATCH] feat(kvp): add diagnostics layer over KvpPoolStore --- libazureinit-kvp/Cargo.toml | 2 +- libazureinit-kvp/src/cli.rs | 231 +++++++++- libazureinit-kvp/src/diagnostics.rs | 631 ++++++++++++++++++++++++++ libazureinit-kvp/src/error.rs | 7 + libazureinit-kvp/src/lib.rs | 6 + libazureinit-kvp/tests/cli.rs | 198 ++++++++ libazureinit-kvp/tests/diagnostics.rs | 234 ++++++++++ 7 files changed, 1306 insertions(+), 3 deletions(-) create mode 100644 libazureinit-kvp/src/diagnostics.rs create mode 100644 libazureinit-kvp/tests/diagnostics.rs diff --git a/libazureinit-kvp/Cargo.toml b/libazureinit-kvp/Cargo.toml index d58f862d..cd0bb03c 100644 --- a/libazureinit-kvp/Cargo.toml +++ b/libazureinit-kvp/Cargo.toml @@ -15,7 +15,7 @@ csv = "1" libc = "0.2" serde_json = "1.0.96" tracing = "0.1.40" -uuid = "1.3" +uuid = { version = "1.3", features = ["v4"] } [dev-dependencies] rstest = { version = "0.26", default-features = false } diff --git a/libazureinit-kvp/src/cli.rs b/libazureinit-kvp/src/cli.rs index fa347cc7..08d3cc83 100644 --- a/libazureinit-kvp/src/cli.rs +++ b/libazureinit-kvp/src/cli.rs @@ -10,8 +10,8 @@ use clap::{Parser, Subcommand, ValueEnum}; use serde_json::json; use crate::{ - write_report, KvpError, KvpPool, KvpPoolStore, PoolMode, - ProvisioningReport, ReportPpsType, + write_report, DiagnosticEvent, DiagnosticRecord, DiagnosticsKvp, KvpError, + KvpPool, KvpPoolStore, PoolMode, ProvisioningReport, ReportPpsType, }; const EXIT_OK: u8 = 0; @@ -176,6 +176,50 @@ enum Command { #[arg(long, value_parser = parse_supporting_data)] supporting_data: Vec, }, + /// Inspect decoded diagnostic events, reassembling chunked records. + Diag { + #[command(subcommand)] + command: DiagCommand, + }, +} + +/// Diagnostics subcommands operating on the reassembled event view. +#[derive(Subcommand, Debug)] +enum DiagCommand { + /// Print every record, reassembling chunked events. Raw records + /// (e.g. PROVISIONING_REPORT) are hidden unless --include-raw. + Dump { + /// Also print raw (non-event) records. + #[arg(long)] + include_raw: bool, + }, + /// Print decoded events (oldest first), optionally filtered. + Events { + /// Only show events at this level (error, warn, info, debug, trace). + #[arg(long)] + level: Option, + /// Only show events whose name contains this substring. + #[arg(long)] + name: Option, + }, + /// Print the last COUNT events (default 20). + Tail { + /// Number of trailing events to print. + #[arg(short = 'n', long = "count", default_value_t = 20)] + count: usize, + }, + /// Remove this layer's events, leaving raw records intact. + Clear { + /// VM identifier whose events to remove. + #[arg(long)] + vm_id: String, + /// Event key prefix whose events to remove. + #[arg(long)] + event_prefix: String, + /// Confirm the removal (required). + #[arg(long)] + yes: bool, + }, } #[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] @@ -267,6 +311,7 @@ fn dispatch(cli: Cli, stdout: &mut W) -> Result { documentation_url, supporting_data, ), + Command::Diag { command } => diag(&store, stdout, command, output), } } @@ -425,6 +470,188 @@ fn is_stale( Ok(if stale { EXIT_OK } else { EXIT_NOT_FOUND }) } +fn diag( + store: &KvpPoolStore, + stdout: &mut W, + command: DiagCommand, + output: OutputMode, +) -> Result { + match command { + DiagCommand::Dump { include_raw } => { + diag_dump(store, stdout, include_raw, output) + } + DiagCommand::Events { level, name } => { + diag_events(store, stdout, level, name, None, output) + } + DiagCommand::Tail { count } => { + diag_events(store, stdout, None, None, Some(count), output) + } + DiagCommand::Clear { + vm_id, + event_prefix, + yes, + } => diag_clear(store, stdout, vm_id, event_prefix, yes, output), + } +} + +fn diag_dump( + store: &KvpPoolStore, + stdout: &mut W, + include_raw: bool, + output: OutputMode, +) -> Result { + let diagnostics = DiagnosticsKvp::new(store.clone(), "", ""); + let records: Vec<_> = diagnostics + .records()? + .into_iter() + .filter(|record| { + include_raw || !matches!(record, DiagnosticRecord::Raw { .. }) + }) + .collect(); + + match output { + OutputMode::Text => { + for record in &records { + let line = match record { + DiagnosticRecord::Event { event, chunks } => format!( + "event level={} name={} event_id={} chunks={} \ + message={}", + event.level, + event.name, + event.event_id, + chunks, + event.message + ), + DiagnosticRecord::Raw { key, value } => { + format!("raw key={key} value={value}") + } + DiagnosticRecord::Malformed { key, value, reason } => { + format!( + "malformed key={key} reason={reason} \ + value={value}" + ) + } + }; + writeln!(stdout, "{line}")?; + } + } + OutputMode::Json => { + let array: Vec<_> = records.iter().map(diag_record_json).collect(); + writeln_json(stdout, &serde_json::Value::Array(array))?; + } + } + Ok(EXIT_OK) +} + +fn diag_events( + store: &KvpPoolStore, + stdout: &mut W, + level: Option, + name: Option, + tail: Option, + output: OutputMode, +) -> Result { + let level = level.as_deref().map(parse_level_filter).transpose()?; + + let diagnostics = DiagnosticsKvp::new(store.clone(), "", ""); + let mut events = diagnostics.events()?; + + if let Some(level) = level { + events.retain(|event| event.level == level); + } + if let Some(needle) = name.as_deref() { + events.retain(|event| event.name.contains(needle)); + } + if let Some(count) = tail { + let excess = events.len().saturating_sub(count); + events.drain(..excess); + } + + match output { + OutputMode::Text => { + for event in &events { + let line = format!( + "event level={} name={} event_id={} message={}", + event.level, event.name, event.event_id, event.message + ); + writeln!(stdout, "{line}")?; + } + } + OutputMode::Json => { + let array: Vec<_> = events.iter().map(diag_event_json).collect(); + writeln_json(stdout, &serde_json::Value::Array(array))?; + } + } + Ok(EXIT_OK) +} + +fn diag_clear( + store: &KvpPoolStore, + stdout: &mut W, + vm_id: String, + event_prefix: String, + yes: bool, + output: OutputMode, +) -> Result { + if !yes { + return Err(CliError::Usage( + "refusing to clear diagnostic events without --yes".to_string(), + )); + } + let diagnostics = DiagnosticsKvp::new(store.clone(), vm_id, event_prefix); + diagnostics.clear()?; + match output { + OutputMode::Text => writeln!(stdout, "cleared")?, + OutputMode::Json => writeln_json(stdout, &json!({ "cleared": true }))?, + } + Ok(EXIT_OK) +} + +/// Parse a `--level` filter argument into a [`tracing::Level`]. +fn parse_level_filter(level: &str) -> Result { + level.parse::().map_err(|_| { + CliError::Usage(format!( + "invalid level '{level}' (expected error, warn, info, debug, \ + or trace)" + )) + }) +} + +/// Render a [`DiagnosticRecord`] as a JSON object. +fn diag_record_json(record: &DiagnosticRecord) -> serde_json::Value { + match record { + DiagnosticRecord::Event { event, chunks } => { + let mut value = diag_event_json(event); + if let serde_json::Value::Object(map) = &mut value { + map.insert("chunks".to_string(), json!(chunks)); + } + value + } + DiagnosticRecord::Raw { key, value } => json!({ + "kind": "raw", + "key": key, + "value": value, + }), + DiagnosticRecord::Malformed { key, value, reason } => json!({ + "kind": "malformed", + "key": key, + "value": value, + "reason": reason, + }), + } +} + +/// Render a [`DiagnosticEvent`] as a JSON object (without chunk count). +fn diag_event_json(event: &DiagnosticEvent) -> serde_json::Value { + json!({ + "kind": "event", + "level": event.level.to_string(), + "name": event.name, + "event_id": event.event_id, + "message": event.message, + }) +} + fn report_success( store: &KvpPoolStore, vm_id: Option, diff --git a/libazureinit-kvp/src/diagnostics.rs b/libazureinit-kvp/src/diagnostics.rs new file mode 100644 index 00000000..e71f74bb --- /dev/null +++ b/libazureinit-kvp/src/diagnostics.rs @@ -0,0 +1,631 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Typed diagnostics layer over the raw +//! [`KvpPoolStore`](crate::KvpPoolStore) key/value API. +//! +//! Where [`KvpPoolStore`](crate::KvpPoolStore) treats keys and values as +//! opaque bytes, [`DiagnosticsKvp`] understands the telemetry +//! conventions azure-init writes into the guest pool: +//! +//! - **Event keys** encode structured metadata as a five-segment, +//! pipe-delimited string +//! (`||||`). +//! - **Chunking**: a single KVP record caps the value at the store's +//! per-record limit (see [`MAX_CHUNK_BYTES`] for the safe-mode value). +//! Longer messages are split at UTF-8 codepoint boundaries into +//! multiple records that share one key, written atomically under a +//! single lock, and reassembled on read. +//! - **Classification**: [`records`](DiagnosticsKvp::records) sorts every +//! stored record into a [`DiagnosticRecord`] — a reassembled +//! [`DiagnosticEvent`], an unstructured [`Raw`](DiagnosticRecord::Raw) +//! record such as `PROVISIONING_REPORT`, or a +//! [`Malformed`](DiagnosticRecord::Malformed) event key. +//! +//! This module is policy only: all locking, size enforcement, and +//! on-disk encoding stay in [`KvpPoolStore`](crate::KvpPoolStore). +//! +//! # Example +//! +//! ``` +//! use libazureinit_kvp::{ +//! DiagnosticEvent, DiagnosticsKvp, KvpPool, KvpPoolStore, PoolMode, +//! MAX_CHUNK_BYTES, +//! }; +//! use tracing::Level; +//! +//! # fn main() -> Result<(), libazureinit_kvp::KvpError> { +//! let dir = std::env::temp_dir() +//! .join(format!("libazureinit-kvp-doc-{}", std::process::id())); +//! std::fs::create_dir_all(&dir)?; +//! let store = KvpPoolStore::new_in(KvpPool::Guest, &dir, PoolMode::Safe)?; +//! store.clear()?; +//! +//! let diagnostics = +//! DiagnosticsKvp::new(store, "vm-1234", "azure-init-doc"); +//! +//! // A short event lands in a single record. +//! diagnostics.emit(&DiagnosticEvent::new( +//! Level::INFO, +//! "user:create_user", +//! "Creating user azureuser", +//! ))?; +//! +//! // A long message is split across records that share one key and is +//! // reassembled transparently on read. +//! let long = "x".repeat(MAX_CHUNK_BYTES * 2 + 10); +//! diagnostics +//! .emit(&DiagnosticEvent::new(Level::DEBUG, "config:dump", &long))?; +//! +//! let events = diagnostics.events()?; +//! assert_eq!(events.len(), 2); +//! assert_eq!(events[1].message.len(), MAX_CHUNK_BYTES * 2 + 10); +//! +//! # std::fs::remove_dir_all(&dir).ok(); +//! # Ok(()) +//! # } +//! ``` + +use tracing::Level; +use uuid::Uuid; + +use crate::{KvpError, KvpPoolStore}; + +/// Maximum number of value bytes per KVP record under +/// [`PoolMode::Safe`](crate::PoolMode::Safe). +/// +/// This matches a safe-mode store's +/// [`KvpPoolStore::max_value_size`](crate::KvpPoolStore::max_value_size). +/// [`DiagnosticsKvp::emit`] splits messages longer than the store's +/// actual limit, so an [`Unsafe`](crate::PoolMode::Unsafe) store uses +/// its larger capacity; this constant is the conservative reference +/// value used throughout the diagnostics conventions. +pub const MAX_CHUNK_BYTES: usize = 1022; + +/// Delimiter separating the segments of a diagnostic event key. +const EVENT_KEY_DELIMITER: char = '|'; + +/// Format a diagnostic event key as its `|`-delimited on-disk string: +/// `||||`. +/// +/// [`classify_key`] is the inverse. For example: +/// +/// ```text +/// azure-init-0.1.0|3f2504e0-4f89-41d3-9a0c-0305e82c3301|INFO|user:create_user|8f3e9c4a-... +/// ``` +fn format_event_key( + prefix: &str, + vm_id: &str, + level: Level, + name: &str, + event_id: &str, +) -> String { + let d = EVENT_KEY_DELIMITER; + format!("{prefix}{d}{vm_id}{d}{level}{d}{name}{d}{event_id}") +} + +/// Outcome of inspecting a raw pool key. +enum KeyClass<'a> { + /// The key is a well-formed event key. + Event { + prefix: &'a str, + vm_id: &'a str, + level: Level, + name: &'a str, + event_id: &'a str, + }, + /// The key has five segments but is not a valid event. + Malformed { reason: String }, + /// The key is not an event key (e.g. `PROVISIONING_REPORT`). + Raw, +} + +/// Classify a raw pool key without allocating. +fn classify_key(key: &str) -> KeyClass<'_> { + let mut segments = key.split(EVENT_KEY_DELIMITER); + + // `str::split` always yields at least one element. + let prefix = segments.next().unwrap_or_default(); + let (Some(vm_id), Some(level), Some(name), Some(event_id)) = ( + segments.next(), + segments.next(), + segments.next(), + segments.next(), + ) else { + return KeyClass::Raw; + }; + if segments.next().is_some() { + // More than five segments: a `|` leaked into a field. + return KeyClass::Raw; + } + + match level.parse::() { + Ok(level) => KeyClass::Event { + prefix, + vm_id, + level, + name, + event_id, + }, + Err(_) => KeyClass::Malformed { + reason: format!("unrecognized level {level:?}"), + }, + } +} + +/// Split `value` into pieces of at most `max_bytes` bytes each, always +/// at UTF-8 codepoint boundaries. +/// +/// An empty input yields a single empty chunk so callers still write one +/// record. A codepoint wider than `max_bytes` (only possible for tiny +/// `max_bytes`, never for [`MAX_CHUNK_BYTES`]) is emitted whole so the +/// split always makes progress. +fn chunk_at_char_boundary(value: &str, max_bytes: usize) -> Vec<&str> { + debug_assert!(max_bytes > 0, "max_bytes must be positive"); + if value.is_empty() { + return vec![""]; + } + + let mut chunks = Vec::new(); + let mut start = 0; + while start < value.len() { + if value.len() - start <= max_bytes { + chunks.push(&value[start..]); + break; + } + + // Walk back from the byte limit to the nearest codepoint boundary. + let mut end = start + max_bytes; + while end > start && !value.is_char_boundary(end) { + end -= 1; + } + if end == start { + // One codepoint spans the whole window; take it whole. + end = start + max_bytes + 1; + while end < value.len() && !value.is_char_boundary(end) { + end += 1; + } + } + + chunks.push(&value[start..end]); + start = end; + } + chunks +} + +/// Reject the `|` key delimiter in an event field so the formatted key +/// round-trips through [`classify_key`]. +fn reject_delimiter(field: &'static str, value: &str) -> Result<(), KvpError> { + if value.contains(EVENT_KEY_DELIMITER) { + return Err(KvpError::EventFieldContainsDelimiter { field }); + } + Ok(()) +} + +/// A single diagnostic event. +/// +/// Construct one with [`new`](Self::new) (which generates a fresh +/// `event_id`) and hand it to [`DiagnosticsKvp::emit`]; events read back +/// via [`DiagnosticsKvp::records`] carry the same fields, decoded from +/// the pool. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DiagnosticEvent { + /// Severity of the event. + pub level: Level, + /// Formatted event name, e.g. `user:create_user`. + pub name: String, + /// Per-emit identifier. [`new`](Self::new) generates a UUIDv4; + /// every chunk of one emitted event shares this value. + pub event_id: String, + /// Literal value bytes written to the pool. The diagnostics layer + /// imposes no format on this string. + pub message: String, +} + +impl DiagnosticEvent { + /// Create an event with a freshly generated `event_id` (UUIDv4). + pub fn new( + level: Level, + name: impl Into, + message: impl Into, + ) -> Self { + Self { + level, + name: name.into(), + event_id: Uuid::new_v4().to_string(), + message: message.into(), + } + } +} + +/// A single record read back from the pool and classified by +/// [`DiagnosticsKvp::records`]. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DiagnosticRecord { + /// A reassembled diagnostic event. + Event { + /// The decoded event. + event: DiagnosticEvent, + /// Number of on-disk records the value spanned (1 when short). + chunks: usize, + }, + /// An unstructured record whose key is not an event key, such as + /// `PROVISIONING_REPORT`. + Raw { + /// The record key. + key: String, + /// The reassembled record value. + value: String, + }, + /// A record whose key has five segments but is not a valid event + /// (for example, an unrecognized level). + Malformed { + /// The record key. + key: String, + /// The reassembled record value. + value: String, + /// Why the key failed to parse as an event. + reason: String, + }, +} + +/// A typed diagnostics view over a [`KvpPoolStore`]. +/// +/// Owns the event-key `prefix` and `vm_id` used to format and +/// [`scope`](Self::clear) this layer's events. See the +/// [module documentation](self) for the on-disk format. +#[derive(Clone, Debug)] +pub struct DiagnosticsKvp { + store: KvpPoolStore, + vm_id: String, + event_prefix: String, +} + +impl DiagnosticsKvp { + /// Wrap `store` with the `vm_id` and `event_prefix` stamped into + /// this layer's event keys. + pub fn new( + store: KvpPoolStore, + vm_id: impl Into, + event_prefix: impl Into, + ) -> Self { + Self { + store, + vm_id: vm_id.into(), + event_prefix: event_prefix.into(), + } + } + + /// The underlying store. + pub fn store(&self) -> &KvpPoolStore { + &self.store + } + + /// The VM identifier stamped into event keys. + pub fn vm_id(&self) -> &str { + &self.vm_id + } + + /// The prefix stamped into event keys. + pub fn event_prefix(&self) -> &str { + &self.event_prefix + } + + /// Write `event`. + /// + /// Messages longer than the store's per-record value limit are split + /// at UTF-8 codepoint boundaries and written as multiple records that + /// share one key, atomically under a single lock via + /// [`KvpPoolStore::append_multiple`]. + /// + /// Returns [`KvpError::EventFieldContainsDelimiter`] if the + /// `event_prefix`, `vm_id`, `name`, or `event_id` contains the `|` + /// key delimiter, which would make the key ambiguous to + /// [`records`](Self::records). + pub fn emit(&self, event: &DiagnosticEvent) -> Result<(), KvpError> { + reject_delimiter("event_prefix", &self.event_prefix)?; + reject_delimiter("vm_id", &self.vm_id)?; + reject_delimiter("name", &event.name)?; + reject_delimiter("event_id", &event.event_id)?; + + let key = format_event_key( + &self.event_prefix, + &self.vm_id, + event.level, + &event.name, + &event.event_id, + ); + + self.write_chunked(&key, &event.message) + } + + /// Split `value` at the store's per-record limit and append every + /// chunk under the same `key` in one atomic batch. + fn write_chunked(&self, key: &str, value: &str) -> Result<(), KvpError> { + let chunks = chunk_at_char_boundary(value, self.store.max_value_size()); + self.store + .append_multiple(chunks.into_iter().map(|chunk| (key, chunk))) + } + + /// Read every record, reassembling chunked events and classifying + /// each into a [`DiagnosticRecord`]. + /// + /// Records are returned in on-disk order. Consecutive records that + /// share a key are one event; because [`emit`](Self::emit) writes an + /// event's chunks contiguously under a single lock, reassembly is + /// correct even under concurrent writers. + pub fn records(&self) -> Result, KvpError> { + Ok(reassemble(self.store.dump()?)) + } + + /// Read back only the records that decode as diagnostic events, in + /// on-disk order. + pub fn events(&self) -> Result, KvpError> { + Ok(self + .records()? + .into_iter() + .filter_map(|record| match record { + DiagnosticRecord::Event { event, .. } => Some(event), + DiagnosticRecord::Raw { .. } + | DiagnosticRecord::Malformed { .. } => None, + }) + .collect()) + } + + /// Remove this layer's events — keys that parse as an event key + /// whose `prefix` and `vm_id` match this instance — under a single + /// lock. Raw records such as `PROVISIONING_REPORT` are left intact. + pub fn clear(&self) -> Result<(), KvpError> { + let keys: Vec = self + .store + .dump()? + .into_iter() + .filter_map(|(key, _)| { + let is_mine = matches!( + classify_key(&key), + KeyClass::Event { prefix, vm_id, .. } + if prefix == self.event_prefix + && vm_id == self.vm_id + ); + is_mine.then_some(key) + }) + .collect(); + self.store.delete_multiple(keys)?; + Ok(()) + } +} + +/// Group consecutive same-key records from [`KvpPoolStore::dump`] and +/// classify each group into a [`DiagnosticRecord`]. +fn reassemble(dumped: Vec<(String, String)>) -> Vec { + let mut records = Vec::new(); + let mut dumped = dumped.into_iter().peekable(); + + while let Some((key, value)) = dumped.next() { + let mut message = value; + let mut chunks = 1; + while dumped.peek().is_some_and(|(next, _)| *next == key) { + let (_, next_value) = dumped.next().expect("peeked value exists"); + message.push_str(&next_value); + chunks += 1; + } + records.push(classify_record(key, message, chunks)); + } + + records +} + +/// Turn one reassembled key/value group into a [`DiagnosticRecord`]. +fn classify_record( + key: String, + value: String, + chunks: usize, +) -> DiagnosticRecord { + match classify_key(&key) { + KeyClass::Event { + level, + name, + event_id, + .. + } => DiagnosticRecord::Event { + event: DiagnosticEvent { + level, + name: name.to_string(), + event_id: event_id.to_string(), + message: value, + }, + chunks, + }, + KeyClass::Malformed { reason } => { + DiagnosticRecord::Malformed { key, value, reason } + } + KeyClass::Raw => DiagnosticRecord::Raw { key, value }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + const PREFIX: &str = "azure-init-0.1.0"; + const VM_ID: &str = "3f2504e0-4f89-41d3-9a0c-0305e82c3301"; + const EVENT_ID: &str = "8f3e9c4a-1b2c-4d5e-9f01-234567890abc"; + + #[test] + fn event_key_formats_and_classifies() { + let formatted = format_event_key( + PREFIX, + VM_ID, + Level::INFO, + "user:create_user", + EVENT_ID, + ); + assert_eq!( + formatted, + format!("{PREFIX}|{VM_ID}|INFO|user:create_user|{EVENT_ID}") + ); + assert!(matches!( + classify_key(&formatted), + KeyClass::Event { + prefix, + vm_id, + level, + name, + event_id, + } if prefix == PREFIX + && vm_id == VM_ID + && level == Level::INFO + && name == "user:create_user" + && event_id == EVENT_ID + )); + } + + #[test] + fn classify_round_trips_every_level() { + for expected in [ + Level::ERROR, + Level::WARN, + Level::INFO, + Level::DEBUG, + Level::TRACE, + ] { + let key = format_event_key( + PREFIX, + VM_ID, + expected, + "span:event", + EVENT_ID, + ); + assert!(matches!( + classify_key(&key), + KeyClass::Event { level, .. } if level == expected + )); + } + } + + /// Map a key to its [`KeyClass`] discriminant for table-driven tests. + fn class_of(key: &str) -> &'static str { + match classify_key(key) { + KeyClass::Event { .. } => "event", + KeyClass::Malformed { .. } => "malformed", + KeyClass::Raw => "raw", + } + } + + #[rstest] + #[case::event("p|vm|INFO|name|id", "event")] + #[case::raw_single_segment("PROVISIONING_REPORT", "raw")] + #[case::raw_too_few_segments("a|b|INFO|c", "raw")] + #[case::raw_too_many_segments("a|b|INFO|c|d|e", "raw")] + #[case::malformed_bad_level("p|vm|NOTALEVEL|name|id", "malformed")] + #[case::malformed_other_level("p|vm|NOPE|name|id", "malformed")] + fn classify_key_categorizes(#[case] key: &str, #[case] expected: &str) { + assert_eq!(class_of(key), expected); + } + + #[rstest] + #[case::empty("", 4, vec![""])] + #[case::shorter_than_max("abc", 8, vec!["abc"])] + #[case::exact_multiple("abcdef", 2, vec!["ab", "cd", "ef"])] + #[case::ascii_remainder("abcde", 2, vec!["ab", "cd", "e"])] + #[case::two_byte_boundary("aéb", 2, vec!["a", "é", "b"])] + #[case::oversized_three_byte("€", 1, vec!["€"])] + #[case::oversized_repeated("€€", 1, vec!["€", "€"])] + fn chunk_splits_at_utf8_boundaries( + #[case] input: &str, + #[case] max_bytes: usize, + #[case] expected: Vec<&str>, + ) { + assert_eq!(chunk_at_char_boundary(input, max_bytes), expected); + } + + #[test] + fn chunk_reassembles_multibyte_payload() { + let payload = "🚀".repeat(100); + let chunks = chunk_at_char_boundary(&payload, 7); + assert!(chunks.iter().all(|chunk| chunk.len() <= 7)); + assert_eq!(chunks.concat(), payload); + } + + #[test] + fn diagnostic_event_new_generates_uuid_event_id() { + let event = DiagnosticEvent::new(Level::INFO, "span:name", "message"); + assert!(Uuid::parse_str(&event.event_id).is_ok()); + } + + #[test] + fn reject_delimiter_flags_pipe() { + assert!(reject_delimiter("name", "no pipe here").is_ok()); + let err = reject_delimiter("name", "has|pipe").unwrap_err(); + assert!(matches!( + err, + KvpError::EventFieldContainsDelimiter { field: "name" } + )); + assert_eq!( + err.to_string(), + "event key field 'name' must not contain '|'" + ); + } + + #[test] + fn reassemble_groups_chunks_and_classifies() { + let key = format_event_key( + PREFIX, + VM_ID, + Level::INFO, + "config:dump", + EVENT_ID, + ); + + let dumped = vec![ + (key.clone(), "part-one/".to_string()), + (key.clone(), "part-two".to_string()), + ( + "PROVISIONING_REPORT".to_string(), + "result=success".to_string(), + ), + ("p|vm|NOPE|name|id".to_string(), "junk".to_string()), + ]; + + let records = reassemble(dumped); + assert_eq!(records.len(), 3); + + assert_eq!( + records[0], + DiagnosticRecord::Event { + event: DiagnosticEvent { + level: Level::INFO, + name: "config:dump".to_string(), + event_id: EVENT_ID.to_string(), + message: "part-one/part-two".to_string(), + }, + chunks: 2, + } + ); + assert!(matches!(&records[1], DiagnosticRecord::Raw { key, .. } + if key == "PROVISIONING_REPORT")); + assert!(matches!(&records[2], DiagnosticRecord::Malformed { .. })); + } + + #[test] + fn reassemble_keeps_distinct_adjacent_keys_separate() { + let make = |event_id: &str| { + format_event_key(PREFIX, VM_ID, Level::INFO, "span:name", event_id) + }; + let dumped = vec![ + (make("id-1"), "first".to_string()), + (make("id-2"), "second".to_string()), + ]; + let records = reassemble(dumped); + assert_eq!(records.len(), 2); + assert!(matches!( + &records[0], + DiagnosticRecord::Event { chunks: 1, .. } + )); + assert!(matches!( + &records[1], + DiagnosticRecord::Event { chunks: 1, .. } + )); + } +} diff --git a/libazureinit-kvp/src/error.rs b/libazureinit-kvp/src/error.rs index 02b9ec9c..71b75abb 100644 --- a/libazureinit-kvp/src/error.rs +++ b/libazureinit-kvp/src/error.rs @@ -11,6 +11,10 @@ pub enum KvpError { EmptyKey, /// An underlying I/O error. Io(io::Error), + /// An event key field (`event_prefix`, `vm_id`, `name`, or + /// `event_id`) contained the `|` delimiter, which would make the + /// formatted event key ambiguous to parse back. + EventFieldContainsDelimiter { field: &'static str }, /// The key contains a null byte, which is incompatible with the /// on-disk format (null-padded fixed-width fields). KeyContainsNull, @@ -29,6 +33,9 @@ impl fmt::Display for KvpError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::EmptyKey => write!(f, "KVP key must not be empty"), + Self::EventFieldContainsDelimiter { field } => { + write!(f, "event key field '{field}' must not contain '|'") + } Self::Io(e) => write!(f, "{e}"), Self::KeyContainsNull => { write!(f, "KVP key must not contain null bytes") diff --git a/libazureinit-kvp/src/lib.rs b/libazureinit-kvp/src/lib.rs index f17b66f7..2e94b14a 100644 --- a/libazureinit-kvp/src/lib.rs +++ b/libazureinit-kvp/src/lib.rs @@ -9,14 +9,20 @@ //! - [`ProvisioningReport`]: structured provisioning health report that //! is persisted as the single `PROVISIONING_REPORT` record with //! [`write_report`]. +//! - [`DiagnosticsKvp`]: typed view over [`KvpPoolStore`] that formats, +//! chunks, and reassembles azure-init diagnostic events. mod cli; +mod diagnostics; mod error; mod report; mod store; mod vm_id; pub use cli::run; +pub use diagnostics::{ + DiagnosticEvent, DiagnosticRecord, DiagnosticsKvp, MAX_CHUNK_BYTES, +}; pub use error::KvpError; pub use report::{ write_report, ProvisioningReport, ReportPpsType, PROVISIONING_REPORT_KEY, diff --git a/libazureinit-kvp/tests/cli.rs b/libazureinit-kvp/tests/cli.rs index ef485314..60416526 100644 --- a/libazureinit-kvp/tests/cli.rs +++ b/libazureinit-kvp/tests/cli.rs @@ -298,3 +298,201 @@ fn report_failure_rejects_invalid_supporting_data() { .unwrap() .contains("key=value")); } + +#[test] +fn diag_dump_json_reassembles_and_classifies() { + let dir = TempDir::new().unwrap(); + // A two-chunk event (same key repeated) plus a raw record. + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "azure-init-x|vm|INFO|a:b|id1", "one/"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "azure-init-x|vm|INFO|a:b|id1", "two"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "PROVISIONING_REPORT", "result=success"], + ))); + + let out = assert_success(kvp(&with_dir(&dir, &["--json", "diag", "dump"]))); + assert!(out.contains("\"kind\":\"event\"")); + assert!(out.contains("\"chunks\":2")); + assert!(out.contains("\"message\":\"one/two\"")); + // Raw records are hidden without --include-raw. + assert!(!out.contains("PROVISIONING_REPORT")); + + let out_raw = assert_success(kvp(&with_dir( + &dir, + &["--json", "diag", "dump", "--include-raw"], + ))); + assert!(out_raw.contains("\"kind\":\"raw\"")); + assert!(out_raw.contains("PROVISIONING_REPORT")); +} + +#[test] +fn diag_events_filters_by_level() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|INFO|a:b|i1", "info-msg"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|ERROR|c:d|i2", "err-msg"], + ))); + + let out = assert_success(kvp(&with_dir( + &dir, + &["diag", "events", "--level", "error"], + ))); + assert!(out.contains("err-msg")); + assert!(!out.contains("info-msg")); +} + +#[test] +fn diag_clear_requires_confirmation_then_scopes_removal() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|INFO|a:b|i1", "msg"], + ))); + + // Without --yes the command refuses and exits with a usage error. + let refused = kvp(&with_dir( + &dir, + &["diag", "clear", "--vm-id", "vm", "--event-prefix", "p"], + )); + assert_eq!(refused.status.code(), Some(2)); + + // With --yes the matching event is removed. + assert_success(kvp(&with_dir( + &dir, + &[ + "diag", + "clear", + "--vm-id", + "vm", + "--event-prefix", + "p", + "--yes", + ], + ))); + let out = assert_success(kvp(&with_dir(&dir, &["diag", "dump"]))); + assert!(out.trim().is_empty()); +} + +#[test] +fn diag_dump_text_renders_all_record_kinds() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|INFO|a:b|id1", "one/"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|INFO|a:b|id1", "two"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "PROVISIONING_REPORT", "result=success"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|NOPE|c:d|id2", "junk"], + ))); + + let out = assert_success(kvp(&with_dir( + &dir, + &["diag", "dump", "--include-raw"], + ))); + assert!(out.contains( + "event level=INFO name=a:b event_id=id1 chunks=2 message=one/two" + )); + assert!(out.contains("raw key=PROVISIONING_REPORT value=result=success")); + assert!(out.contains("malformed key=p|vm|NOPE|c:d|id2")); + assert!(out.contains("value=junk")); +} + +#[test] +fn diag_tail_limits_to_last_events() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|INFO|a:b|i1", "first"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|INFO|c:d|i2", "second"], + ))); + + let out = + assert_success(kvp(&with_dir(&dir, &["diag", "tail", "-n", "1"]))); + assert!(out.contains("second")); + assert!(!out.contains("first")); +} + +#[test] +fn diag_events_filters_by_name_substring() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|INFO|user:add|i1", "u"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|INFO|ssh:key|i2", "s"], + ))); + + let out = assert_success(kvp(&with_dir( + &dir, + &["diag", "events", "--name", "ssh"], + ))); + assert!(out.contains("ssh:key")); + assert!(!out.contains("user:add")); +} + +#[test] +fn diag_events_rejects_invalid_level() { + let dir = TempDir::new().unwrap(); + let output = kvp(&with_dir(&dir, &["diag", "events", "--level", "bogus"])); + assert_eq!(output.status.code(), Some(2)); +} + +#[test] +fn diag_json_covers_malformed_events_and_clear() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|INFO|a:b|i1", "hello"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "p|vm|NOPE|c:d|i2", "junk"], + ))); + + let dump = + assert_success(kvp(&with_dir(&dir, &["--json", "diag", "dump"]))); + assert!(dump.contains("\"kind\":\"event\"")); + assert!(dump.contains("\"kind\":\"malformed\"")); + assert!(dump.contains("\"reason\":")); + + let events = + assert_success(kvp(&with_dir(&dir, &["--json", "diag", "events"]))); + assert!(events.contains("\"message\":\"hello\"")); + let cleared = assert_success(kvp(&with_dir( + &dir, + &[ + "--json", + "diag", + "clear", + "--vm-id", + "vm", + "--event-prefix", + "p", + "--yes", + ], + ))); + assert!(cleared.contains("\"cleared\":true")); +} diff --git a/libazureinit-kvp/tests/diagnostics.rs b/libazureinit-kvp/tests/diagnostics.rs new file mode 100644 index 00000000..d846aab0 --- /dev/null +++ b/libazureinit-kvp/tests/diagnostics.rs @@ -0,0 +1,234 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Integration tests for the [`DiagnosticsKvp`] layer: emit/read +//! round-trips, chunk reassembly, classification, scoped clearing, and +//! the concurrent-write atomicity guarantee that keeps chunked events +//! from interleaving. + +use std::thread; + +use libazureinit_kvp::{ + DiagnosticEvent, DiagnosticRecord, DiagnosticsKvp, KvpPool, KvpPoolStore, + PoolMode, MAX_CHUNK_BYTES, +}; +use tempfile::TempDir; +use tracing::Level; + +const PREFIX: &str = "azure-init-test"; +const VM_ID: &str = "vm-abc"; + +fn diagnostics(dir: &TempDir) -> DiagnosticsKvp { + let store = + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), PoolMode::Safe) + .unwrap(); + DiagnosticsKvp::new(store, VM_ID, PREFIX) +} + +#[test] +fn short_event_round_trips_as_single_record() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + assert_eq!(diag.vm_id(), VM_ID); + assert_eq!(diag.event_prefix(), PREFIX); + + let event = + DiagnosticEvent::new(Level::INFO, "user:create_user", "created"); + diag.emit(&event).unwrap(); + + assert_eq!(diag.store().dump().unwrap().len(), 1); + + let records = diag.records().unwrap(); + assert_eq!(records.len(), 1); + match &records[0] { + DiagnosticRecord::Event { + event: decoded, + chunks, + } => { + assert_eq!(*chunks, 1); + assert_eq!(decoded.level, Level::INFO); + assert_eq!(decoded.name, "user:create_user"); + assert_eq!(decoded.event_id, event.event_id); + assert_eq!(decoded.message, "created"); + } + other => panic!("expected event, got {other:?}"), + } +} + +#[test] +fn long_event_splits_across_records_and_reassembles() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + let message = "x".repeat(MAX_CHUNK_BYTES * 3 + 50); + let event = DiagnosticEvent::new(Level::DEBUG, "config:dump", &message); + diag.emit(&event).unwrap(); + + // Split across four records that all share one key (split keys). + let dumped = diag.store().dump().unwrap(); + assert_eq!(dumped.len(), 4); + let key = &dumped[0].0; + assert!(dumped.iter().all(|(k, _)| k == key)); + + let records = diag.records().unwrap(); + assert_eq!(records.len(), 1); + match &records[0] { + DiagnosticRecord::Event { + event: decoded, + chunks, + } => { + assert_eq!(*chunks, 4); + assert_eq!(decoded.message, message); + } + other => panic!("expected event, got {other:?}"), + } +} + +#[test] +fn injected_malformed_key_is_classified() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + // Five segments but an unrecognized level. + diag.store() + .append(&format!("{PREFIX}|{VM_ID}|NOPE|bad:level|id"), "junk") + .unwrap(); + + let records = diag.records().unwrap(); + assert_eq!(records.len(), 1); + assert!(matches!( + &records[0], + DiagnosticRecord::Malformed { reason, .. } if reason.contains("NOPE") + )); +} + +#[test] +fn mixed_records_round_trip_together() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + diag.emit(&DiagnosticEvent::new(Level::INFO, "a:b", "short")) + .unwrap(); + diag.emit(&DiagnosticEvent::new( + Level::WARN, + "c:d", + "y".repeat(MAX_CHUNK_BYTES + 5), + )) + .unwrap(); + diag.store() + .append("PROVISIONING_REPORT", "result=success") + .unwrap(); + diag.store() + .append(&format!("{PREFIX}|{VM_ID}|NOPE|e:f|id"), "junk") + .unwrap(); + + let records = diag.records().unwrap(); + // Two events + one raw + one malformed. + assert_eq!(records.len(), 4); + assert_eq!(diag.events().unwrap().len(), 2); +} + +#[test] +fn clear_removes_events_but_keeps_raw() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + diag.emit(&DiagnosticEvent::new(Level::INFO, "a:b", "e1")) + .unwrap(); + diag.emit(&DiagnosticEvent::new( + Level::DEBUG, + "c:d", + "z".repeat(MAX_CHUNK_BYTES * 2), + )) + .unwrap(); + diag.store() + .append("PROVISIONING_REPORT", "result=success") + .unwrap(); + + diag.clear().unwrap(); + + let records = diag.records().unwrap(); + assert_eq!(records.len(), 1); + assert!(matches!( + &records[0], + DiagnosticRecord::Raw { key, .. } if key == "PROVISIONING_REPORT" + )); + assert!(diag.events().unwrap().is_empty()); +} + +#[test] +fn clear_is_scoped_to_matching_prefix_and_vm_id() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + diag.emit(&DiagnosticEvent::new(Level::INFO, "a:b", "mine")) + .unwrap(); + // An event from a different agent/VM must survive clear(). + diag.store() + .append("other-agent|other-vm|INFO|x:y|id", "theirs") + .unwrap(); + + diag.clear().unwrap(); + + let events = diag.events().unwrap(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].message, "theirs"); +} + +#[test] +fn emit_rejects_delimiter_in_event_fields() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + // A pipe in the name would produce an ambiguous six-segment key. + let event = DiagnosticEvent::new(Level::INFO, "a|b", "msg"); + assert!(diag.emit(&event).is_err()); + // Nothing was written. + assert!(diag.store().dump().unwrap().is_empty()); +} + +#[test] +fn concurrent_multichunk_emits_reassemble_without_interleaving() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + const THREADS: usize = 5; + const PER_THREAD: usize = 8; + // Force three chunks per event. + let len = MAX_CHUNK_BYTES * 2 + 7; + + let handles: Vec<_> = (0..THREADS) + .map(|t| { + let diag = diag.clone(); + let marker = (b'a' + t as u8) as char; + thread::spawn(move || { + for _ in 0..PER_THREAD { + let message = marker.to_string().repeat(len); + let event = DiagnosticEvent::new( + Level::INFO, + format!("thread:{marker}"), + message, + ); + diag.emit(&event).unwrap(); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + + let events = diag.events().unwrap(); + // If any event's chunks had been split by an interleaving writer, the + // key would appear as multiple groups and the count would be wrong. + assert_eq!(events.len(), THREADS * PER_THREAD); + for event in &events { + // Each message is homogeneous and full length: chunks stayed + // contiguous on disk. + assert_eq!(event.message.len(), len); + let first = event.message.chars().next().unwrap(); + assert!(event.message.chars().all(|c| c == first)); + assert_eq!(event.name, format!("thread:{first}")); + } +}