diff --git a/crates/trident/src/grpc_client/mod.rs b/crates/trident/src/grpc_client/mod.rs index 9eae37298..46d78f7ca 100644 --- a/crates/trident/src/grpc_client/mod.rs +++ b/crates/trident/src/grpc_client/mod.rs @@ -20,7 +20,7 @@ mod tridentclient; use error::TridentClientError; use tridentclient::{RebootHandling, TridentClient}; -pub fn client_main(args: &ClientArgs) -> ExitCode { +pub fn client_main(args: &ClientArgs, before_reboot: impl FnOnce()) -> ExitCode { // Start the Tokio runtime let Ok(runtime) = Builder::new_multi_thread().enable_all().build() else { error!("Failed to create Tokio runtime"); @@ -78,9 +78,10 @@ pub fn client_main(args: &ClientArgs) -> ExitCode { error!("Client failed: {:?}", e); return TridentExitCodes::Failed.into(); } + // Daemon-owned reboot may drop pending client telemetry. Ok(ExitKind::Done) => {} Ok(ExitKind::NeedsReboot) => { - if let Err(e) = crate::request_reboot_with_wait() { + if let Err(e) = crate::request_reboot_with_wait_and_flush(before_reboot) { error!("Failed to reboot: {e:?}"); return TridentExitCodes::RebootUnsuccessful.into(); } diff --git a/crates/trident/src/lib.rs b/crates/trident/src/lib.rs index 06a0e8888..a51f6509e 100644 --- a/crates/trident/src/lib.rs +++ b/crates/trident/src/lib.rs @@ -59,6 +59,7 @@ pub use crate::{ }, grpc_client::client_main, logging::{ + appinsights::AppInsightsSender, background_log::BackgroundLog, background_uploader::{BackgroundUploadHandle, BackgroundUploader}, journald::JournaldTracing, @@ -69,10 +70,11 @@ pub use crate::{ run_with_captured_operation, save_reboot_operation, take_reboot_operation, OperationSource, }, + telemetry_uploader::TelemetryUploader, tracestream::TraceStream, }, orchestrate::OrchestratorConnection, - reboot::request_reboot_with_wait, + reboot::{request_reboot_with_wait, request_reboot_with_wait_and_flush}, server::server_main, }; diff --git a/crates/trident/src/logging/appinsights.rs b/crates/trident/src/logging/appinsights.rs new file mode 100644 index 000000000..257a1a96f --- /dev/null +++ b/crates/trident/src/logging/appinsights.rs @@ -0,0 +1,1172 @@ +//! Best-effort tracing sink that forwards Trident's metric/span tracing +//! events to Azure Monitor / Application Insights. +//! +//! This intentionally does not depend on the OpenTelemetry SDK or an +//! Application Insights client crate. It follows the same minimal approach +//! as [`super::tracestream::TraceSender`]: parse the Application Insights +//! *connection string* (`InstrumentationKey=;IngestionEndpoint=;...`) +//! ourselves and build the raw Application Insights `EventData` envelope. +//! +//! Sending is delegated to [`super::telemetry_uploader`], separate from the +//! log-forwarding uploader used by [`super::logstream::Logstream`]. +//! `send_event` only enqueues the envelope; the actual POST runs on a dedicated +//! thread, and failures are logged but never allowed to affect servicing. + +use crate::telemetry::fields; +use std::{ + collections::BTreeMap, + io::Read, + sync::{Arc, RwLock}, + time::{Duration, Instant}, +}; + +use anyhow::Context; +use log::trace; +use serde::{ + de::{IgnoredAny, SeqAccess, Visitor}, + Deserialize, Deserializer, +}; +use serde_json::{json, Value}; +use tracing::{ + field::{Field, Visit}, + span, Event, Subscriber, +}; +use tracing_subscriber::{layer::Layer, registry::LookupSpan}; +use url::Url; + +use super::{ + telemetry_uploader::TelemetryUploadHandle, + tracestream::{merge_operation_context, PLATFORM_INFO}, + upload_core::ResponseValidator, +}; +use crate::TRIDENT_VERSION; + +/// Default Application Insights ingestion endpoint, used when the connection +/// string does not specify one explicitly. +const DEFAULT_INGESTION_ENDPOINT: &str = "https://dc.services.visualstudio.com"; + +/// Per-request total timeout, enforced by the background uploader. Telemetry +/// must never meaningfully delay Trident's actual work. +const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); + +/// `Content-Type` for the Application Insights ingestion request. +const CONTENT_TYPE_JSON: &str = "application/json"; + +/// A parsed Application Insights connection string. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ConnParts { + /// Ingestion endpoint, trailing slash stripped (e.g. + /// `https://region.in.applicationinsights.azure.com`). + pub ingestion_endpoint: String, + /// Instrumentation key. + pub instrumentation_key: String, +} + +impl ConnParts { + /// The `POST` target: `${ingestion_endpoint}/v2/track`. + fn track_url(&self) -> Option { + Url::parse(&format!( + "{}/v2/track", + self.ingestion_endpoint.trim_end_matches('/') + )) + .ok() + } +} + +/// Parse an Application Insights connection string of the form +/// `InstrumentationKey=;IngestionEndpoint=https://...;...`. Returns `None` +/// if the string is empty, unparsable, or missing an instrumentation key. +/// +/// If `IngestionEndpoint` is not given explicitly, the endpoint is derived +/// from the sovereign-cloud `EndpointSuffix`/`Location` fields when present +/// (e.g. `EndpointSuffix=applicationinsights.azure.cn;Location=chinaeast2` -> +/// `https://chinaeast2.dc.applicationinsights.azure.cn`), matching the +/// Azure Monitor SDKs' documented connection-string format, which always +/// uses the `dc` ingestion prefix (with or without `Location`) -- not `in`. +/// Only if none of `IngestionEndpoint`/`EndpointSuffix` are present does +/// this fall back to the public Application Insights endpoint -- a +/// sovereign-cloud string must never be silently redirected to the public +/// endpoint. +pub(crate) fn parse_connection_string(s: &str) -> Option { + let mut instrumentation_key: Option = None; + let mut ingestion_endpoint: Option = None; + let mut endpoint_suffix: Option = None; + let mut location: Option = None; + + for part in s.split(';') { + let part = part.trim(); + if part.is_empty() { + continue; + } + let (key, value) = part.split_once('=')?; + match key.trim().to_ascii_lowercase().as_str() { + "" => return None, + "instrumentationkey" => instrumentation_key = Some(value.trim().to_string()), + "ingestionendpoint" => { + let endpoint = value.trim().trim_end_matches('/'); + if endpoint.is_empty() { + return None; + } + ingestion_endpoint = Some(endpoint.to_string()); + } + "endpointsuffix" => { + let suffix = value.trim().trim_matches('/'); + if suffix.is_empty() { + return None; + } + endpoint_suffix = Some(suffix.to_string()); + } + "location" => location = Some(value.trim().to_string()), + _ => {} + } + } + + let instrumentation_key = instrumentation_key.filter(|k| !k.is_empty())?; + + let ingestion_endpoint = match ingestion_endpoint { + Some(explicit) => explicit, + None => match endpoint_suffix { + // Sovereign-cloud form: derive the ingestion endpoint from + // EndpointSuffix (+ optional Location), rather than assuming the + // public endpoint. The ingestion prefix is always `dc` + // (matching the public endpoint's own + // `dc.services.visualstudio.com` shape), with or without + // Location -- not `in`, which the Azure Monitor SDKs never + // produce and which would target a nonstandard/nonexistent + // host for sovereign clouds. + Some(suffix) => match location.filter(|l| !l.is_empty()) { + Some(location) => format!("https://{location}.dc.{suffix}"), + None => format!("https://dc.{suffix}"), + }, + None => DEFAULT_INGESTION_ENDPOINT.to_string(), + }, + }; + + Some(ConnParts { + ingestion_endpoint, + instrumentation_key, + }) +} + +/// A visitor that records the fields of a tracing event/span as a +/// `BTreeMap`, mirroring [`super::tracestream::TraceEntryVisitor`]. +#[derive(Default)] +struct FieldVisitor { + fields: BTreeMap, +} + +impl Visit for FieldVisitor { + fn record_i64(&mut self, field: &Field, value: i64) { + self.fields.insert(field.name().to_string(), json!(value)); + } + + fn record_f64(&mut self, field: &Field, value: f64) { + self.fields.insert(field.name().to_string(), json!(value)); + } + + fn record_u64(&mut self, field: &Field, value: u64) { + self.fields.insert(field.name().to_string(), json!(value)); + } + + fn record_bool(&mut self, field: &Field, value: bool) { + self.fields.insert(field.name().to_string(), json!(value)); + } + + fn record_str(&mut self, field: &Field, value: &str) { + self.fields.insert(field.name().to_string(), json!(value)); + } + + fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) { + self.fields + .insert(field.name().to_string(), json!(format!("{value:?}"))); + } +} + +/// Timestamp recorded when a span is entered, used to compute execution time +/// on exit. +struct SpanStart(Instant); + +/// Renders a JSON value as a string, since Application Insights `EventData` +/// properties are a `Map`. +fn stringify(value: &Value) -> String { + match value { + Value::String(s) => s.clone(), + other => other.to_string(), + } +} + +/// A `tracing_subscriber::Layer` that forwards Trident's metric events and +/// instrumented spans to Application Insights, best-effort. +/// +/// Only constructed when a connection string was compiled into the binary +/// (see [`crate::AZURE_MONITOR_CONNECTION_STRING`]) *and* telemetry has been +/// enabled via the Agent Configuration file (see +/// [`crate::agentconfig::AgentConfig::telemetry_enabled`]); see +/// [`AppInsightsSender::from_connection_string`]. +pub struct AppInsightsSender { + instrumentation_key: String, + track_url: Url, + uploader: TelemetryUploadHandle, + /// The same persistent, per-host installation ID handle used by + /// `TraceStream`/`TraceSender` (see `TraceStream::installation_id_handle`), + /// so Application Insights events can be correlated back to a specific + /// host installation the same way tracestream metrics already are. + installation_id: Arc>>, + /// The same servicing ID handle used by `TraceStream`/`TraceSender` + /// (see `TraceStream::servicing_id_handle`), correlating events across + /// a stage/finalize/commit sequence for one servicing operation. + servicing_id: Arc>>, +} + +impl AppInsightsSender { + /// Build a sender from an Application Insights connection string. + /// Returns `None` if the string is empty, fails to parse, the + /// ingestion endpoint does not form a valid URL, or that URL's scheme + /// is not `https`. + /// + /// The events sent through this sender include host identifiers (see + /// [`super::tracestream::PLATFORM_INFO`]), so a non-HTTPS endpoint -- + /// e.g. from a build-time typo/misconfiguration -- is rejected rather + /// than silently sending opted-in host telemetry in cleartext. Azure + /// Monitor ingestion endpoints require HTTPS. + pub fn from_connection_string( + connection_string: &str, + uploader: TelemetryUploadHandle, + installation_id: Arc>>, + servicing_id: Arc>>, + ) -> Option { + let parts = parse_connection_string(connection_string)?; + match parts.track_url() { + Some(url) if url.scheme() == "https" => {} + _ => { + trace!( + "Application Insights ingestion endpoint '{}' is not HTTPS, disabling telemetry", + parts.ingestion_endpoint + ); + return None; + } + } + Self::from_parts(parts, uploader, installation_id, servicing_id) + } + + fn from_parts( + parts: ConnParts, + uploader: TelemetryUploadHandle, + installation_id: Arc>>, + servicing_id: Arc>>, + ) -> Option { + let track_url = parts.track_url()?; + Some(Self { + instrumentation_key: parts.instrumentation_key, + track_url, + uploader, + installation_id, + servicing_id, + }) + } + + /// Build and enqueue an Application Insights `EventData` envelope for + /// the background uploader to send, best-effort. This only serializes + /// the envelope and hands it off to the uploader's channel, so it never + /// blocks on network I/O. A serialization failure or a closed uploader + /// is logged here at `trace` level; a later network error or non-2xx + /// response is logged by the background uploader itself (at `error` + /// level, same as any other background upload). Either way the failure + /// is otherwise ignored -- it can never affect servicing outcomes. + fn send_event(&self, name: &str, mut properties: BTreeMap) { + properties.insert(fields::TRIDENT_VERSION.to_string(), json!(TRIDENT_VERSION)); + for (key, value) in PLATFORM_INFO.iter() { + properties.insert(key.clone(), json!(stringify(value))); + } + if let Ok(installation_id) = self.installation_id.read() { + if let Some(installation_id) = installation_id.as_ref() { + properties + .entry(fields::INSTALLATION_ID.to_string()) + .or_insert_with(|| json!(installation_id)); + } + } + if let Ok(servicing_id) = self.servicing_id.read() { + if let Some(servicing_id) = servicing_id.as_ref() { + properties + .entry(fields::SERVICING_ID.to_string()) + .or_insert_with(|| json!(servicing_id)); + } + } + // operation_id/command/source enrichment is + // shared with the local metrics-file sink -- see + // `merge_operation_context`'s doc comment for why. + merge_operation_context(name, &mut properties); + + let string_properties: BTreeMap = properties + .into_iter() + .map(|(key, value)| (key, stringify(&value))) + .collect(); + + let now = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true); + let envelope = json!({ + "name": format!( + "Microsoft.ApplicationInsights.{}.Event", + self.instrumentation_key.replace('-', "") + ), + "time": now, + "iKey": self.instrumentation_key, + "tags": { "ai.internal.sdkVersion": format!("trident:{TRIDENT_VERSION}") }, + "data": { + "baseType": "EventData", + "baseData": { + "ver": 2, + "name": name, + "properties": string_properties, + } + } + }); + + let body = match serde_json::to_vec(&envelope) { + Ok(b) => b, + Err(e) => { + trace!("Failed to serialize Application Insights event: {e}"); + return; + } + }; + + if let Err(e) = self.uploader.upload_with_validator( + &self.track_url, + body, + REQUEST_TIMEOUT, + Some(CONTENT_TYPE_JSON), + Some(ResponseValidator { + body_status: reqwest::StatusCode::PARTIAL_CONTENT, + validate: validate_track_response, + }), + ) { + trace!("Failed to enqueue Application Insights event: {e}"); + } + } +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct TrackResponse { + items_received: u64, + items_accepted: u64, + #[serde(default, deserialize_with = "deserialize_first_track_error")] + errors: Option, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct TrackError { + status_code: Option, +} + +fn deserialize_first_track_error<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + struct FirstError; + + impl<'de> Visitor<'de> for FirstError { + type Value = Option; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("an array of ingestion errors or null") + } + + fn visit_unit(self) -> Result { + Ok(None) + } + + fn visit_seq(self, mut sequence: A) -> Result + where + A: SeqAccess<'de>, + { + let first = sequence.next_element()?; + while sequence.next_element::()?.is_some() {} + Ok(first) + } + } + + deserializer.deserialize_any(FirstError) +} + +/// Response validator for the Application Insights `/v2/track` endpoint +/// (see [`TelemetryUploadHandle::upload_with_validator`]). Each request here +/// carries exactly one envelope, so only `itemsReceived == 1 && itemsAccepted == 1` +/// counts as success; anything else is treated like a failed upload. +fn validate_track_response(reader: &mut dyn Read) -> Result<(), anyhow::Error> { + let response: TrackResponse = serde_json::from_reader(reader) + .context("Failed to parse Application Insights partial-success response body")?; + + if response.items_received != 1 || response.items_accepted != 1 { + let rejection = response + .errors + .and_then(|error| error.status_code) + .map(|status| format!("; first rejection status {status}")) + .unwrap_or_default(); + anyhow::bail!( + "Application Insights accepted only {} of {} submitted items \ + (expected exactly 1 of 1 for this single-envelope request){}", + response.items_accepted, + response.items_received, + rejection, + ); + } + + Ok(()) +} + +/// The `Layer` implementation mirrors +/// [`super::tracestream::TraceSender`]'s event/span handling, but renders an +/// Application Insights `EventData` envelope instead of Trident's own +/// metrics-file format. +impl Layer for AppInsightsSender +where + S: Subscriber + for<'span> LookupSpan<'span>, +{ + fn enabled( + &self, + metadata: &tracing::Metadata<'_>, + _ctx: tracing_subscriber::layer::Context<'_, S>, + ) -> bool { + metadata.level() <= &tracing::Level::INFO + } + + fn on_event(&self, event: &Event<'_>, _ctx: tracing_subscriber::layer::Context<'_, S>) { + let mut visitor = FieldVisitor::default(); + event.record(&mut visitor); + + let Some(metric_name) = visitor + .fields + .get(fields::METRIC_NAME) + .and_then(|v| v.as_str()) + .map(str::to_string) + else { + // Not a metric event (e.g. a plain log line); nothing to forward. + return; + }; + + let properties: BTreeMap = visitor + .fields + .into_iter() + .filter(|(key, _)| key != fields::METRIC_NAME) + .collect(); + + self.send_event(&metric_name, properties); + } + + fn on_new_span( + &self, + attrs: &span::Attributes<'_>, + id: &span::Id, + ctx: tracing_subscriber::layer::Context<'_, S>, + ) { + if let Some(span) = ctx.span(id) { + let mut visitor = FieldVisitor::default(); + attrs.record(&mut visitor); + span.extensions_mut().insert(visitor); + } + } + + fn on_enter(&self, id: &span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) { + if let Some(span) = ctx.span(id) { + span.extensions_mut().insert(SpanStart(Instant::now())); + } + } + + fn on_exit(&self, id: &span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) { + let Some(span) = ctx.span(id) else { + return; + }; + let Some(SpanStart(start)) = span.extensions_mut().remove::() else { + return; + }; + let execution_time = start.elapsed().as_secs_f64(); + let Some(mut fields) = span + .extensions() + .get::() + .map(|visitor| visitor.fields.clone()) + else { + return; + }; + + fields.insert(fields::EXECUTION_TIME.to_string(), json!(execution_time)); + + self.send_event(span.name(), fields); + } + + fn on_record( + &self, + id: &span::Id, + values: &span::Record<'_>, + ctx: tracing_subscriber::layer::Context<'_, S>, + ) { + if let Some(span) = ctx.span(id) { + if let Some(visitor) = span.extensions_mut().get_mut::() { + values.record(visitor); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_reentered_span_preserves_fields() { + use super::super::telemetry_uploader::TelemetryUploader; + use tracing_subscriber::{layer::SubscriberExt, Registry}; + + let mut server = mockito::Server::new(); + let uploader = TelemetryUploader::new(8).unwrap(); + let sender = AppInsightsSender::from_parts( + ConnParts { + ingestion_endpoint: server.url(), + instrumentation_key: "test-key".to_string(), + }, + uploader.get_handle().unwrap(), + Arc::new(RwLock::new(None)), + Arc::new(RwLock::new(None)), + ) + .unwrap(); + let mocks: Vec<_> = ["first", "second"] + .into_iter() + .map(|phase| { + server + .mock("POST", "/v2/track") + .match_body(mockito::Matcher::PartialJson(json!({ + "data": {"baseData": { + "name": "reentered_operation", + "properties": {"phase": phase, "initial": "kept"} + }} + }))) + .match_request(|request| { + let envelope: Value = + serde_json::from_slice(request.body().unwrap()).unwrap(); + envelope["data"]["baseData"]["properties"]["execution_time"] + .as_str() + .and_then(|value| value.parse::().ok()) + .is_some_and(|value| value >= 0.0) + }) + .with_status(200) + .create() + }) + .collect(); + + tracing::subscriber::with_default(Registry::default().with(sender), || { + let span = + tracing::info_span!("reentered_operation", phase = "first", initial = "kept"); + { + let _entered = span.enter(); + } + span.record("phase", "second"); + { + let _entered = span.enter(); + } + }); + assert!(uploader.flush_with_deadline(Duration::from_secs(5))); + for mock in mocks { + mock.assert(); + } + uploader.shutdown_with_deadline(Duration::from_secs(5)); + } + + #[test] + fn test_startup_events_omit_persisted_ids() { + use super::super::{operation_context, telemetry_uploader::TelemetryUploader}; + use crate::command_kind::CommandKind; + + for command in [ + CommandKind::install_stage(), + CommandKind::install_finalize(), + ] { + for source in [ + operation_context::OperationSource::Cli, + operation_context::OperationSource::Daemon, + operation_context::OperationSource::GrpcClient, + ] { + let mut server = mockito::Server::new(); + let uploader = TelemetryUploader::new(8).unwrap(); + let sender = AppInsightsSender::from_parts( + ConnParts { + ingestion_endpoint: server.url(), + instrumentation_key: "test-key".to_string(), + }, + uploader.get_handle().unwrap(), + Arc::new(RwLock::new(Some("shared-installation".to_string()))), + Arc::new(RwLock::new(Some("previous-servicing".to_string()))), + ) + .unwrap(); + + operation_context::run_with_operation(&command, source, || { + let operation_id = operation_context::current().unwrap().0; + let servicing_id = if command.generates_servicing_id() + && operation_context::should_generate_persistent_ids(source) + { + operation_id.as_str() + } else { + "previous-servicing" + }; + let mut mocks = Vec::new(); + for name in ["command_start", "trident_start", "clean_install_start"] { + let mut expected_properties = json!({ + "operation_id": operation_id, + "command": command.as_str(), + "source": source.as_str(), + "value": "42" + }); + if name == "clean_install_start" { + expected_properties["installation_id"] = json!("shared-installation"); + expected_properties["servicing_id"] = json!(servicing_id); + } + let mock = server + .mock("POST", "/v2/track") + .match_body(mockito::Matcher::PartialJson(json!({ + "data": {"baseData": { + "name": name, + "properties": expected_properties + }} + }))) + .match_request(move |request| { + let envelope: Value = + serde_json::from_slice(request.body().unwrap()).unwrap(); + let properties = &envelope["data"]["baseData"]["properties"]; + name == "clean_install_start" + || (properties.get("installation_id").is_none() + && properties.get("servicing_id").is_none()) + }) + .with_status(200) + .create(); + let mut properties = BTreeMap::from([("value".to_string(), json!(42))]); + if name != "clean_install_start" { + properties.insert("installation_id".to_string(), json!("explicit-id")); + properties.insert("servicing_id".to_string(), json!("explicit-id")); + } + mocks.push(mock); + sender.send_event(name, properties); + } + assert!(uploader.flush_with_deadline(Duration::from_secs(5))); + for mock in mocks { + mock.assert(); + } + }); + uploader.shutdown_with_deadline(Duration::from_secs(5)); + } + } + } + + #[test] + fn test_span_timing_is_captured_before_following_layers() { + use std::sync::atomic::{AtomicBool, Ordering}; + use tracing_subscriber::{layer::SubscriberExt, Registry}; + + struct ObserveExit(Arc); + impl Layer for ObserveExit { + fn on_exit( + &self, + id: &span::Id, + ctx: tracing_subscriber::layer::Context<'_, Registry>, + ) { + let span = ctx.span(id).unwrap(); + assert!(span.extensions().get::().is_none()); + self.0.store(true, Ordering::SeqCst); + } + } + + let observed = Arc::new(AtomicBool::new(false)); + let sender = AppInsightsSender::from_connection_string( + "InstrumentationKey=k;IngestionEndpoint=https://region.example/", + TelemetryUploadHandle::new_mock(), + Arc::new(RwLock::new(None)), + Arc::new(RwLock::new(None)), + ) + .unwrap(); + let mut layers: Vec + Send + Sync>> = + vec![Box::new(ObserveExit(observed.clone()))]; + layers.insert(0, Box::new(sender)); + tracing::subscriber::with_default(Registry::default().with(layers), || { + let span = tracing::info_span!("timed_operation"); + let _entered = span.enter(); + }); + assert!(observed.load(Ordering::SeqCst)); + } + + #[test] + fn test_validate_track_response_206_all_items_accepted() { + let body = br#"{"itemsReceived":1,"itemsAccepted":1}"#; + assert!(validate_track_response(&mut body.as_slice()).is_ok()); + } + + #[test] + fn test_validate_track_response_206_item_rejected() { + let body = br#"{"itemsReceived":1,"itemsAccepted":0}"#; + let err = validate_track_response(&mut body.as_slice()) + .expect_err("a rejected single-envelope request should be treated as a failure"); + assert!(err.to_string().contains("accepted only 0 of 1")); + } + + #[test] + fn test_validate_track_response_reports_first_rejection_status() { + for status in [400, 429, 500] { + let body = json!({ + "itemsReceived": 1, + "itemsAccepted": 0, + "errors": [ + {"index": 0, "statusCode": status, "message": "rejected"}, + {"statusCode": 503} + ] + }) + .to_string(); + let error = validate_track_response(&mut body.as_bytes()).unwrap_err(); + let message = error.to_string(); + assert!(message.contains("accepted only 0 of 1")); + assert!(message.contains(&format!("first rejection status {status}"))); + assert!(!message.contains("503")); + } + } + + #[test] + fn test_validate_track_response_streams_rejection_details() { + let prefix = + br#"{"itemsReceived":1,"itemsAccepted":0,"errors":[{"statusCode":429,"message":""#; + let suffix = br#""},{"statusCode":500}]}"#; + let mut reader = prefix + .as_slice() + .chain(std::io::repeat(b'x').take(1024 * 1024)) + .chain(suffix.as_slice()); + + let error = validate_track_response(&mut reader).unwrap_err(); + let message = error.to_string(); + assert!(message.contains("first rejection status 429")); + assert!(message.len() < 200); + assert!(!message.contains("500")); + } + + #[test] + fn test_validate_track_response_empty_rejection_details() { + for errors in [json!([]), Value::Null, json!([{}])] { + for accepted in [0, 1] { + let body = json!({ + "itemsReceived": 1, + "itemsAccepted": accepted, + "errors": errors + }) + .to_string(); + let result = validate_track_response(&mut body.as_bytes()); + assert_eq!(result.is_ok(), accepted == 1); + if let Err(error) = result { + assert!(error.to_string().contains("accepted only 0 of 1")); + } + } + } + } + + #[test] + fn test_validate_track_response_206_zero_zero_is_not_success() { + // A malformed/impossible response for a single-envelope request -- + // mere equality between the two counts (0 == 0) must NOT be + // mistaken for success. + let body = br#"{"itemsReceived":0,"itemsAccepted":0}"#; + let err = validate_track_response(&mut body.as_slice()) + .expect_err("itemsReceived=0/itemsAccepted=0 must not be treated as success"); + assert!(err.to_string().contains("accepted only 0 of 0")); + } + + #[test] + fn test_validate_track_response_206_impossible_counts_above_one() { + // Also impossible for a single-envelope request -- equality + // alone (2 == 2) must not be mistaken for success either. + let body = br#"{"itemsReceived":2,"itemsAccepted":2}"#; + let err = validate_track_response(&mut body.as_slice()) + .expect_err("itemsReceived=2/itemsAccepted=2 is impossible for a 1-item request"); + assert!(err.to_string().contains("accepted only 2 of 2")); + } + + #[test] + fn test_validate_track_response_206_unparsable_body() { + let err = validate_track_response(&mut b"not json".as_slice()).expect_err( + "an unparsable 206 body should be treated as a failure, not silently accepted", + ); + assert!(err.to_string().contains("Failed to parse")); + } + + #[test] + fn test_validate_track_response_streams_large_ignored_fields() { + let prefix = br#"{"errors":[{"message":""#; + let suffix = br#""}],"itemsAccepted":1,"itemsReceived":1}"#; + let mut reader = prefix + .as_slice() + .chain(std::io::repeat(b'x').take(1024 * 1024)) + .chain(suffix.as_slice()); + + validate_track_response(&mut reader).unwrap(); + } + + #[test] + fn test_validate_track_response_rejects_invalid_fields_and_trailing_data() { + for body in [ + r#"{"itemsReceived":1}"#, + r#"{"itemsReceived":1,"itemsAccepted":"1"}"#, + r#"{"itemsReceived":1,"itemsAccepted":-1}"#, + r#"{"itemsReceived":1,"itemsAccepted":1,"itemsAccepted":0}"#, + r#"{"itemsReceived":1,"itemsAccepted":1} trailing"#, + r#"{"itemsReceived":1,"itemsAccepted":1,"errors":["#, + ] { + assert!( + validate_track_response(&mut body.as_bytes()).is_err(), + "accepted invalid response: {body}" + ); + } + } + + #[test] + fn test_validate_track_response_propagates_body_read_failure() { + struct FailingReader; + + impl Read for FailingReader { + fn read(&mut self, _: &mut [u8]) -> std::io::Result { + Err(std::io::Error::other("response read failed")) + } + } + + let prefix = br#"{"itemsReceived":1,"itemsAccepted":1}"#; + let mut reader = prefix.as_slice().chain(FailingReader); + + let err = validate_track_response(&mut reader).unwrap_err(); + assert!(format!("{err:#}").contains("response read failed")); + } + + #[test] + fn test_parse_connection_string_full() { + let parts = parse_connection_string( + "InstrumentationKey=abc123;IngestionEndpoint=https://region.example/;LiveEndpoint=https://live.example/", + ) + .expect("should parse"); + assert_eq!(parts.instrumentation_key, "abc123"); + assert_eq!(parts.ingestion_endpoint, "https://region.example"); + assert_eq!( + parts.track_url(), + Some(Url::parse("https://region.example/v2/track").unwrap()) + ); + } + + #[test] + fn test_parse_connection_string_missing_endpoint_uses_default() { + let parts = parse_connection_string("InstrumentationKey=abc123").expect("should parse"); + assert_eq!(parts.instrumentation_key, "abc123"); + assert_eq!(parts.ingestion_endpoint, DEFAULT_INGESTION_ENDPOINT); + assert_eq!( + parts.track_url(), + Some(Url::parse(&format!("{DEFAULT_INGESTION_ENDPOINT}/v2/track")).unwrap()) + ); + } + + #[test] + fn test_parse_connection_string_rejects_explicitly_empty_endpoints() { + for connection in [ + "InstrumentationKey=k;IngestionEndpoint=;", + "InstrumentationKey=k;IngestionEndpoint= ;", + "InstrumentationKey=k;EndpointSuffix=;", + "InstrumentationKey=k;EndpointSuffix= ;", + "InstrumentationKey=k;IngestionEndpoint=;EndpointSuffix=applicationinsights.azure.cn", + "InstrumentationKey=k;IngestionEndpoint=https://region.example/;EndpointSuffix=", + "InstrumentationKey=k;IngestionEndpoint=;IngestionEndpoint=https://region.example/", + "InstrumentationKey=k;EndpointSuffix= ;EndpointSuffix=applicationinsights.azure.cn", + "InstrumentationKey=k;IngestionEndpoint=/;IngestionEndpoint=https://region.example/", + "InstrumentationKey=k;EndpointSuffix=/;EndpointSuffix=applicationinsights.azure.cn", + ] { + assert!( + parse_connection_string(connection).is_none(), + "accepted {connection}" + ); + } + } + + #[test] + fn test_parse_connection_string_rejects_malformed_fields() { + for connection in [ + "InstrumentationKey=k;EndpointSuffix", + "InstrumentationKey=k;IngestionEndpoint", + "InstrumentationKey=k; EndpointSuffix ;", + "InstrumentationKey=k;UnknownField", + "InstrumentationKey=k;=value", + "InstrumentationKey=k; =value", + "EndpointSuffix;InstrumentationKey=k", + ] { + assert!( + parse_connection_string(connection).is_none(), + "accepted {connection}" + ); + } + } + + #[test] + fn test_parse_connection_string_preserves_valid_fields() { + for (connection, expected) in [ + ("InstrumentationKey=k", DEFAULT_INGESTION_ENDPOINT), + (" ; InstrumentationKey = k ; ; ", DEFAULT_INGESTION_ENDPOINT), + ("InstrumentationKey=k;UnknownField=value", DEFAULT_INGESTION_ENDPOINT), + ( + " instrumentationkey = k ; ingestionendpoint = https://region.example/ ; ", + "https://region.example", + ), + ( + "InstrumentationKey=k;EndpointSuffix=applicationinsights.azure.cn;Location=chinaeast2", + "https://chinaeast2.dc.applicationinsights.azure.cn", + ), + ( + "InstrumentationKey=k;IngestionEndpoint=https://first.example/;IngestionEndpoint=https://last.example/", + "https://last.example", + ), + ] { + let parts = parse_connection_string(connection).expect(connection); + assert_eq!(parts.ingestion_endpoint, expected); + } + } + + #[test] + /// Sovereign-cloud connection strings that specify `EndpointSuffix` (and + /// optionally `Location`) instead of `IngestionEndpoint` must derive the + /// matching sovereign ingestion endpoint, not silently fall back to the + /// public one. + fn test_parse_connection_string_endpoint_suffix_with_location() { + let parts = parse_connection_string( + "InstrumentationKey=abc123;EndpointSuffix=applicationinsights.azure.cn;Location=chinaeast2", + ) + .expect("should parse"); + assert_eq!(parts.instrumentation_key, "abc123"); + assert_eq!( + parts.ingestion_endpoint, + "https://chinaeast2.dc.applicationinsights.azure.cn" + ); + } + + #[test] + fn test_parse_connection_string_endpoint_suffix_without_location() { + let parts = parse_connection_string( + "InstrumentationKey=abc123;EndpointSuffix=applicationinsights.azure.cn", + ) + .expect("should parse"); + assert_eq!( + parts.ingestion_endpoint, + "https://dc.applicationinsights.azure.cn" + ); + } + + #[test] + fn test_parse_connection_string_explicit_ingestion_endpoint_wins_over_suffix() { + let parts = parse_connection_string( + "InstrumentationKey=abc123;IngestionEndpoint=https://region.example/;EndpointSuffix=applicationinsights.azure.cn", + ) + .expect("should parse"); + assert_eq!(parts.ingestion_endpoint, "https://region.example"); + } + + #[test] + fn test_parse_connection_string_missing_ikey_is_none() { + assert!(parse_connection_string("IngestionEndpoint=https://region.example/").is_none()); + assert!(parse_connection_string( + "InstrumentationKey=;IngestionEndpoint=https://region.example/" + ) + .is_none()); + } + + #[test] + fn test_parse_connection_string_empty_is_none() { + assert!(parse_connection_string("").is_none()); + } + + #[test] + fn test_from_connection_string_empty_is_none() { + assert!(AppInsightsSender::from_connection_string( + "", + TelemetryUploadHandle::new_mock(), + Arc::new(RwLock::new(None)), + Arc::new(RwLock::new(None)), + ) + .is_none()); + } + + #[test] + fn test_from_connection_string_builds_sender() { + let sender = AppInsightsSender::from_connection_string( + "InstrumentationKey=k;IngestionEndpoint=https://region.example/", + TelemetryUploadHandle::new_mock(), + Arc::new(RwLock::new(None)), + Arc::new(RwLock::new(None)), + ) + .expect("should build sender"); + assert_eq!(sender.instrumentation_key, "k"); + assert_eq!( + sender.track_url, + Url::parse("https://region.example/v2/track").unwrap() + ); + } + + #[test] + /// A non-HTTPS ingestion endpoint (e.g. from a build-time typo/ + /// misconfiguration) must be rejected rather than silently accepted -- + /// events sent through this sender include host identifiers. + fn test_from_connection_string_rejects_non_https_endpoint() { + assert!(AppInsightsSender::from_connection_string( + "InstrumentationKey=k;IngestionEndpoint=http://region.example/", + TelemetryUploadHandle::new_mock(), + Arc::new(RwLock::new(None)), + Arc::new(RwLock::new(None)), + ) + .is_none()); + } + + #[test] + fn test_stringify() { + assert_eq!(stringify(&json!("hello")), "hello"); + assert_eq!(stringify(&json!(42)), "42"); + assert_eq!(stringify(&json!(true)), "true"); + } +} + +#[cfg(feature = "functional-test")] +#[cfg_attr(not(test), allow(unused_imports, dead_code))] +mod functional_test { + use super::*; + use crate::command_kind::CommandKind; + // Only used by this feature-gated module (a plain `cargo check`/`cargo + // test` without `--features functional-test` never compiles this mod, + // which would otherwise make the top-level import unused). + use crate::logging::operation_context; + + use std::{ + io::{Read, Write}, + net::{TcpListener, TcpStream}, + sync::mpsc::channel, + }; + + use pytest_gen::functional_test; + use tracing_subscriber::{filter, layer::SubscriberExt}; + + /// Reads a full HTTP request off `stream`. A single `TcpStream::read` + /// call is not guaranteed to return the entire request (TCP is a byte + /// stream, not message-oriented) -- keep reading until the header + /// terminator has arrived and, per the declared `Content-Length`, the + /// full body has too. + fn read_full_http_request(stream: &mut TcpStream) -> String { + let mut buf: Vec = Vec::new(); + let mut chunk = [0u8; 4096]; + + loop { + match stream.read(&mut chunk) { + Ok(0) => break, // Peer closed the connection. + Ok(n) => buf.extend_from_slice(&chunk[..n]), + Err(_) => break, + } + + let Some(headers_end) = find_subslice(&buf, b"\r\n\r\n") else { + continue; // Headers not fully received yet. + }; + let headers = String::from_utf8_lossy(&buf[..headers_end]); + let content_length: usize = headers + .lines() + .find_map(|line| { + line.to_lowercase() + .strip_prefix("content-length:") + .map(|v| v.trim().to_string()) + }) + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + + let body_end = headers_end + content_length; + if buf.len() >= body_end { + buf.truncate(body_end); + break; + } + } + + String::from_utf8_lossy(&buf).to_string() + } + + fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .position(|window| window == needle) + .map(|pos| pos + needle.len()) + } + + /// Spins up a local TCP listener standing in for the Application + /// Insights ingestion endpoint, and confirms the sender actually posts a + /// well-formed `EventData` envelope to it over the network. + #[functional_test] + fn test_app_insights_sender_posts_event() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + + // command_start (fired by run_with_operation) and test_metric below + // are each posted as their own request, so accept and collect every + // connection the listener sees rather than assuming exactly one. + let (tx, rx) = channel(); + std::thread::spawn(move || { + for stream in listener.incoming() { + let Ok(mut stream) = stream else { break }; + let received = read_full_http_request(&mut stream); + // Send a minimal response so the client's request completes + // cleanly instead of hitting a connection reset. + let _ = stream.write_all( + b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ); + if tx.send(received).is_err() { + break; + } + } + }); + + let uploader = crate::TelemetryUploader::new(16).expect("should build uploader"); + let sender = AppInsightsSender::from_parts( + ConnParts { + ingestion_endpoint: format!("http://{addr}"), + instrumentation_key: "test-key".to_string(), + }, + uploader.get_handle().expect("uploader should be alive"), + Arc::new(RwLock::new(Some("test-installation-id".to_string()))), + Arc::new(RwLock::new(Some("test-servicing-id".to_string()))), + ) + .expect("should build sender") + .with_filter(filter::LevelFilter::INFO); + + let _guard = + tracing::subscriber::set_default(tracing_subscriber::Registry::default().with(sender)); + + // Wrapping in run_with_operation confirms operation_id/command/source + // also reach the outgoing properties, alongside installation_id above. + operation_context::run_with_operation( + &CommandKind::for_test("test_command"), + operation_context::OperationSource::Cli, + || { + tracing::info!(metric_name = "test_metric", value = true); + }, + ); + + // Collect both requests (command_start + test_metric); order between + // them is not guaranteed, so gather everything seen within the + // timeout and assert across the combined traffic. + let mut requests = Vec::new(); + while requests.len() < 2 { + match rx.recv_timeout(std::time::Duration::from_secs(5)) { + Ok(req) => requests.push(req), + Err(e) => panic!("did not receive the expected requests: {e:?}"), + } + } + let combined = requests.join("\n"); + + assert!(combined.contains("POST /v2/track")); + // The envelope's outer "name" strips dashes from the instrumentation + // key (per the Application Insights envelope convention), while + // "iKey" below keeps the original dashed key. + assert!(combined.contains("\"name\":\"Microsoft.ApplicationInsights.testkey.Event\"")); + assert!(combined.contains("\"name\":\"test_metric\"")); + assert!(combined.contains("\"iKey\":\"test-key\"")); + assert!(combined.contains("\"installation_id\":\"test-installation-id\"")); + assert!(combined.contains("\"servicing_id\":\"test-servicing-id\"")); + assert!(combined.contains("\"command\":\"test_command\"")); + assert!(combined.contains("\"operation_id\":")); + assert!(combined.contains("\"source\":\"cli\"")); + } +} diff --git a/crates/trident/src/logging/mod.rs b/crates/trident/src/logging/mod.rs index 06d8d1802..1219de457 100644 --- a/crates/trident/src/logging/mod.rs +++ b/crates/trident/src/logging/mod.rs @@ -1,12 +1,14 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; +pub(super) mod appinsights; pub(super) mod background_log; pub(super) mod background_uploader; pub(super) mod journald; pub(super) mod logfwd; pub(super) mod logstream; pub(super) mod operation_context; +pub(super) mod telemetry_uploader; pub(super) mod tracestream; pub(super) mod upload_core; diff --git a/crates/trident/src/logging/operation_context.rs b/crates/trident/src/logging/operation_context.rs index bbcbec106..60f123b9a 100644 --- a/crates/trident/src/logging/operation_context.rs +++ b/crates/trident/src/logging/operation_context.rs @@ -1,6 +1,7 @@ //! Thread-local "which command is currently executing, under what //! operation ID, and from which of Trident's three entry points" context, -//! so the local telemetry sink ([`super::tracestream::TraceSender`]) can tag every metric/span +//! so telemetry sinks ([`super::tracestream::TraceSender`], +//! [`super::appinsights::AppInsightsSender`]) can tag every metric/span //! fired during a command's execution with `command`/`operation_id`/ //! `source` fields, without every call site (deep in `engine::*`, //! `Trident::*`, etc.) needing to pass them explicitly. diff --git a/crates/trident/src/logging/telemetry_uploader.rs b/crates/trident/src/logging/telemetry_uploader.rs new file mode 100644 index 000000000..5661056d9 --- /dev/null +++ b/crates/trident/src/logging/telemetry_uploader.rs @@ -0,0 +1,505 @@ +use std::{ + collections::VecDeque, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Condvar, Mutex, Weak, + }, + thread::JoinHandle, + time::Duration, +}; + +use anyhow::{bail, Error}; +use log::{debug, error}; +use tokio::sync::Notify; +use url::Url; + +use super::upload_core::{ + self, attempt_upload, join_with_deadline, AttemptOutcome, OriginCooldowns, ResponseValidator, + UploadData, +}; + +/// A bounded, non-blocking FIFO queue for telemetry uploads. +/// Pushing past `capacity` evicts the oldest not-yet-attempted item instead +/// of blocking producers or growing without bound. +/// +/// This is for best-effort telemetry: producers must never block on enqueue, +/// but a merely slow endpoint still needs a backlog cap. +struct TelemetryQueue { + pending: VecDeque, + in_flight: bool, +} + +struct TelemetryRing { + queue: Mutex, + drained: Condvar, + capacity: usize, + /// Woken on each push, and once more by `signal_stop`, so `pop` can + /// wake for either new data or shutdown. + notify: Notify, + /// Set during shutdown. `ring_upload_loop` only uses this to decide + /// whether to keep draining after a failed upload. + stop_requested: AtomicBool, +} + +impl TelemetryRing { + fn new(capacity: usize) -> Arc { + Arc::new(Self { + queue: Mutex::new(TelemetryQueue { + pending: VecDeque::with_capacity(capacity), + in_flight: false, + }), + drained: Condvar::new(), + capacity, + notify: Notify::new(), + stop_requested: AtomicBool::new(false), + }) + } + + /// Enqueues `item`, or refuses it once shutdown has been signaled. + /// The stop check and enqueue share the `queue` lock used by + /// `signal_stop`/`pop`, so `push` cannot report success for an item that + /// shutdown has already made unreachable. + fn push(&self, item: UploadData) -> Result<(), Error> { + let mut queue = self.queue.lock().unwrap(); + if self.stop_requested.load(Ordering::Acquire) { + bail!("Background uploader has been shut down"); + } + if queue.pending.len() >= self.capacity { + queue.pending.pop_front(); + } + queue.pending.push_back(item); + drop(queue); + self.notify.notify_one(); + Ok(()) + } + + /// Waits for and returns the next item, or `None` once shutdown has + /// been requested and the queue has been fully drained. + async fn pop(&self) -> Option { + loop { + { + let mut queue = self.queue.lock().unwrap(); + if let Some(item) = queue.pending.pop_front() { + queue.in_flight = true; + return Some(item); + } + if self.stop_requested.load(Ordering::Acquire) { + return None; + } + } + self.notify.notified().await; + } + } + + fn finish_upload(&self) { + let mut queue = self.queue.lock().unwrap(); + queue.in_flight = false; + self.drained.notify_all(); + } + + fn flush_with_deadline(&self, deadline: Duration) -> bool { + let queue = self.queue.lock().unwrap(); + let (queue, _) = self + .drained + .wait_timeout_while(queue, deadline, |queue| { + !queue.pending.is_empty() || queue.in_flight + }) + .unwrap(); + queue.pending.is_empty() && !queue.in_flight + } + + /// Signals shutdown under the same `queue` lock used by `push`/`pop`; see + /// `push` for the race this prevents. + fn signal_stop(&self) { + let _queue = self.queue.lock().unwrap(); + self.stop_requested.store(true, Ordering::Release); + drop(_queue); + self.notify.notify_one(); + } +} + +/// Background uploader for best-effort Application Insights telemetry. +/// Unlike [`super::background_uploader::BackgroundUploader`], this uses a +/// bounded [`TelemetryRing`] so telemetry cannot build an unbounded backlog. +/// POST/backoff logic is shared through [`super::upload_core`]. +pub struct TelemetryUploader { + inner: Option<(Arc, JoinHandle<()>)>, +} + +impl TelemetryUploader { + /// Creates a new telemetry uploader with a queue capped at `capacity`. + /// Telemetry producers must never block on enqueue, but a slow endpoint + /// still should not be allowed to grow backlog without bound. + pub fn new(capacity: usize) -> Result { + if capacity == 0 { + bail!("telemetry uploader capacity must be greater than zero"); + } + let ring = TelemetryRing::new(capacity); + let handle = upload_core::spawn_uploader_thread(Self::ring_upload_loop(ring.clone()))?; + Ok(Self { + inner: Some((ring, handle)), + }) + } + + /// Gets a handle to send data to the uploader. Returns `None` if the uploader has been shut down. + pub fn get_handle(&self) -> Option { + let (ring, _) = self.inner.as_ref()?; + Some(TelemetryUploadHandle { + ring: Arc::downgrade(ring), + }) + } + + /// Main upload loop for the telemetry ring. + /// During normal operation it behaves like the log uploader loop. After + /// shutdown is requested, the first failed upload stops draining so the + /// bounded shutdown window is not spent retrying an already failing endpoint. + async fn ring_upload_loop(ring: Arc) { + let mut origin_cooldowns = OriginCooldowns::new(); + + while let Some(upload) = ring.pop().await { + let draining = ring.stop_requested.load(Ordering::Acquire); + let outcome = attempt_upload(&mut origin_cooldowns, upload).await; + ring.finish_upload(); + if draining && matches!(outcome, AttemptOutcome::Attempted { succeeded: false }) { + debug!("Telemetry uploader stopping drain after a failed upload during shutdown"); + break; + } + } + + debug!("Telemetry uploader loop has exited"); + } + + /// Waits for queued and in-flight attempts without closing the uploader. + /// Returns false if the deadline expires before it becomes idle. + pub fn flush_with_deadline(&self, deadline: Duration) -> bool { + let Some((ring, _)) = self.inner.as_ref() else { + return true; + }; + ring.flush_with_deadline(deadline) + } + + /// Signals shutdown and waits up to `deadline` for the uploader thread + /// to drain queued work and exit. + /// + /// Call this for telemetry when shutdown must be bounded. If the deadline + /// expires first, the thread is abandoned and this call returns immediately. + pub fn shutdown_with_deadline(mut self, deadline: Duration) { + let Some((ring, handle)) = self.inner.take() else { + return; + }; + ring.signal_stop(); + + match join_with_deadline(handle, deadline) { + Ok(Ok(())) => debug!("Telemetry uploader shut down"), + Ok(Err(e)) => error!("Telemetry uploader thread panicked: {:?}", e), + Err(_) => { + debug!("Telemetry uploader did not shut down within {deadline:?}; abandoning it") + } + } + } +} + +impl Drop for TelemetryUploader { + fn drop(&mut self) { + // Signaling shutdown lets the upload loop exit gracefully on its own. + if let Some((ring, handle)) = self.inner.take() { + ring.signal_stop(); + debug!("Waiting for telemetry uploader to shut down"); + if let Err(e) = handle.join() { + error!("Telemetry uploader thread panicked: {:?}", e); + } + } + } +} + +/// A handle to send data to the telemetry uploader. +#[derive(Clone)] +pub struct TelemetryUploadHandle { + ring: Weak, +} + +impl TelemetryUploadHandle { + /// Sends data to be uploaded in the background. Any 2xx response is + /// treated as success; use [`Self::upload_with_validator`] if the + /// destination's ingestion protocol can reject part of a request + /// while still returning a 2xx status. + pub fn upload( + &self, + url: &Url, + body: impl Into>, + timeout: Duration, + content_type: Option<&'static str>, + ) -> Result<(), Error> { + self.upload_with_validator(url, body, timeout, content_type, None) + } + + /// Same as [`Self::upload`], but with an optional response validator. + /// See `upload_core::UploadData::response_validator` for the contract. + pub(super) fn upload_with_validator( + &self, + url: &Url, + body: impl Into>, + timeout: Duration, + content_type: Option<&'static str>, + response_validator: Option, + ) -> Result<(), Error> { + let data = UploadData { + url: url.clone(), + body: body.into(), + timeout, + content_type, + response_validator, + }; + if let Some(ring) = self.ring.upgrade() { + ring.push(data) + } else { + bail!("Background uploader has been shut down"); + } + } + + /// Creates a new mock handle that does nothing. + #[cfg(test)] + pub fn new_mock() -> Self { + let ring = TelemetryRing::new(0); + ring.signal_stop(); + Self { + ring: Arc::downgrade(&ring), + } + // `ring` itself is dropped here (only a `Weak` is kept), so the + // handle behaves like one whose uploader has already shut down -- + // matching `BackgroundUploadHandle::new_mock`'s "does nothing" + // contract. + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use mockito::Matcher; + + fn init_test_logging() { + let _ = env_logger::builder() + .filter_level(log::LevelFilter::Trace) + .is_test(true) + .try_init(); + } + + fn run_in_runtime(f: impl std::future::Future) { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(f); + } + + fn mock_upload(url: Url, body: &str) -> UploadData { + UploadData { + url, + body: body.as_bytes().to_vec(), + timeout: Duration::from_secs(2), + content_type: None, + response_validator: None, + } + } + + #[test] + fn test_flush_waits_for_in_flight_attempt_without_closing() { + let ring = TelemetryRing::new(2); + ring.push(mock_upload( + Url::parse("http://example.invalid/").unwrap(), + "reboot", + )) + .unwrap(); + run_in_runtime(async { + assert!(ring.pop().await.is_some()); + }); + + let started = std::time::Instant::now(); + assert!(!ring.flush_with_deadline(Duration::from_millis(20))); + assert!(started.elapsed() < Duration::from_secs(1)); + ring.finish_upload(); + assert!(ring.flush_with_deadline(Duration::ZERO)); + ring.push(mock_upload( + Url::parse("http://example.invalid/").unwrap(), + "command_error", + )) + .unwrap(); + assert!(!ring.flush_with_deadline(Duration::ZERO)); + } + + #[test] + fn test_non_closing_flush_allows_later_error_upload() { + let mut server = mockito::Server::new(); + let reboot = server + .mock("POST", "/reboot") + .with_status(200) + .expect(1) + .create(); + let error = server + .mock("POST", "/command-error") + .with_status(200) + .expect(1) + .create(); + let uploader = TelemetryUploader::new(4).unwrap(); + let handle = uploader.get_handle().unwrap(); + let url = Url::parse(&server.url()).unwrap(); + + handle + .upload( + &url.join("reboot").unwrap(), + b"reboot".to_vec(), + Duration::from_secs(2), + None, + ) + .unwrap(); + assert!(uploader.flush_with_deadline(Duration::from_secs(2))); + reboot.assert(); + handle + .upload( + &url.join("command-error").unwrap(), + b"failure".to_vec(), + Duration::from_secs(2), + None, + ) + .unwrap(); + uploader.shutdown_with_deadline(Duration::from_secs(2)); + error.assert(); + } + + #[test] + /// A `TelemetryRing` over capacity evicts the *oldest* not-yet-attempted + /// entry, preserving FIFO order for the survivors, instead of blocking + /// the caller or rejecting the new item. + fn test_telemetry_ring_evicts_oldest_when_over_capacity() { + let ring = TelemetryRing::new(2); + let url = Url::parse("http://example.invalid/").unwrap(); + + ring.push(mock_upload(url.clone(), "first")).unwrap(); + ring.push(mock_upload(url.clone(), "second")).unwrap(); + // Over capacity: "first" should be evicted, not "second". + ring.push(mock_upload(url.clone(), "third")).unwrap(); + + run_in_runtime(async { + let first = ring.pop().await.unwrap(); + assert_eq!(first.body, b"second"); + let second = ring.pop().await.unwrap(); + assert_eq!(second.body, b"third"); + }); + } + + #[test] + /// Before `signal_stop`, `pop()` waits for new items rather than + /// returning `None` -- a ring with nothing queued yet is not the same + /// as a shut-down one. + fn test_telemetry_ring_pop_waits_for_push_before_stop() { + let ring = TelemetryRing::new(4); + let url = Url::parse("http://example.invalid/").unwrap(); + + run_in_runtime(async { + let ring_clone = ring.clone(); + let waiter = tokio::spawn(async move { ring_clone.pop().await }); + + // Give the waiter a moment to actually start waiting before + // pushing, so this also exercises the `Notify` wakeup path + // rather than just the "already queued" fast path. + tokio::time::sleep(Duration::from_millis(20)).await; + ring.push(mock_upload(url, "hello")).unwrap(); + + let item = tokio::time::timeout(Duration::from_secs(1), waiter) + .await + .expect("pop() should return once an item is pushed") + .unwrap(); + assert_eq!(item.unwrap().body, b"hello"); + }); + } + + #[test] + /// Once `signal_stop` is called, `pop()` drains whatever is left and + /// then returns `None` -- it does not wait for further pushes. + fn test_telemetry_ring_pop_returns_none_after_stop_and_drain() { + let ring = TelemetryRing::new(4); + let url = Url::parse("http://example.invalid/").unwrap(); + ring.push(mock_upload(url, "queued-before-stop")).unwrap(); + ring.signal_stop(); + + run_in_runtime(async { + let item = ring.pop().await; + assert_eq!(item.unwrap().body, b"queued-before-stop"); + assert!(ring.pop().await.is_none()); + }); + } + + #[test] + /// `ring_upload_loop`'s core shutdown behavior: while draining after + /// `signal_stop`, the *first* failed upload stops the loop outright, + /// abandoning whatever is left queued, rather than continuing through + /// the whole backlog. + fn test_ring_upload_loop_stops_draining_after_first_failure_once_stopped() { + init_test_logging(); + + let mut server = mockito::Server::new(); + // First item: a bad status the loop will treat as a failure. + let failing = server + .mock("POST", "/fail") + .with_status(500) + .expect(1) + .create(); + // Second item: queued behind the failure, but should never be + // attempted once the loop stops draining after that failure. + let should_not_hit = server + .mock("POST", "/skip") + .with_status(200) + .expect(0) + .create(); + + let ring = TelemetryRing::new(4); + ring.push(mock_upload( + Url::parse(&server.url()).unwrap().join("/fail").unwrap(), + "will-fail", + )) + .unwrap(); + ring.push(mock_upload( + Url::parse(&server.url()).unwrap().join("/skip").unwrap(), + "should-be-abandoned", + )) + .unwrap(); + // Signal stop *before* the loop starts draining, so both queued + // items are already present when draining begins. + ring.signal_stop(); + + run_in_runtime(async { + TelemetryUploader::ring_upload_loop(ring).await; + }); + + failing.assert(); + should_not_hit.assert(); + } + + #[test] + /// End-to-end: `TelemetryUploader::new` behaves like a normal uploader + /// for a single successful upload, exercising the ring path through the + /// public API rather than `TelemetryRing`/`ring_upload_loop` directly. + fn test_telemetry_uploader_new_sends_post_request() { + init_test_logging(); + + let mut server = mockito::Server::new(); + let body = "hello-telemetry-uploader"; + let mock = server + .mock("POST", "/upload") + .match_body(Matcher::Exact(body.to_string())) + .with_status(200) + .expect(1) + .create(); + + let uploader = TelemetryUploader::new(4).unwrap(); + let handle = uploader.get_handle().unwrap(); + let url = Url::parse(&server.url()).unwrap().join("/upload").unwrap(); + handle + .upload(&url, body.as_bytes().to_vec(), Duration::from_secs(2), None) + .unwrap(); + + drop(uploader); + mock.assert(); + } +} diff --git a/crates/trident/src/logging/upload_core.rs b/crates/trident/src/logging/upload_core.rs index bbac51c00..fbd1c8c11 100644 --- a/crates/trident/src/logging/upload_core.rs +++ b/crates/trident/src/logging/upload_core.rs @@ -166,8 +166,7 @@ pub(super) fn backoff_for_failures(consecutive_failures: u32) -> Duration { /// Used by the telemetry ring loop to tell "skipped due to cooldown" from /// an upload that was actually attempted. pub(super) enum AttemptOutcome { - Succeeded, - Failed, + Attempted { succeeded: bool }, SkippedCooldown, } @@ -216,7 +215,7 @@ pub(super) async fn attempt_upload( // future failure starts from the base cooldown rather // than a previously-escalated one. origin_cooldowns.remove(&origin); - AttemptOutcome::Succeeded + AttemptOutcome::Attempted { succeeded: true } } Err(e) => { error!("Background upload failed: {e}"); @@ -246,7 +245,7 @@ pub(super) async fn attempt_upload( } ); - AttemptOutcome::Failed + AttemptOutcome::Attempted { succeeded: false } } } } @@ -388,7 +387,7 @@ mod tests { }); let mut cooldowns = OriginCooldowns::new(); let outcome = attempt_upload(&mut cooldowns, upload).await; - succeeded = matches!(outcome, AttemptOutcome::Succeeded); + succeeded = matches!(outcome, AttemptOutcome::Attempted { succeeded: true }); cooled_down = !cooldowns.is_empty(); }); release_tx.send(()).unwrap(); @@ -432,7 +431,10 @@ mod tests { run_in_runtime(async { let mut cooldowns = OriginCooldowns::new(); let outcome = attempt_upload(&mut cooldowns, mock_upload(url, body)).await; - assert!(matches!(outcome, AttemptOutcome::Succeeded)); + assert!(matches!( + outcome, + AttemptOutcome::Attempted { succeeded: true } + )); }); mock.assert(); } @@ -459,7 +461,10 @@ mod tests { let mut cooldowns = OriginCooldowns::new(); let fail_url = Url::parse(&server.url()).unwrap().join("/fail").unwrap(); let outcome = attempt_upload(&mut cooldowns, mock_upload(fail_url, "will-fail")).await; - assert!(matches!(outcome, AttemptOutcome::Failed)); + assert!(matches!( + outcome, + AttemptOutcome::Attempted { succeeded: false } + )); // Same origin, still within the cooldown window: should be // skipped outright rather than attempted. @@ -493,7 +498,10 @@ mod tests { run_in_runtime(async { let mut cooldowns = OriginCooldowns::new(); let outcome = attempt_upload(&mut cooldowns, mock_upload(url, "redirect-me")).await; - assert!(matches!(outcome, AttemptOutcome::Failed)); + assert!(matches!( + outcome, + AttemptOutcome::Attempted { succeeded: false } + )); }); redirect_mock.assert(); @@ -550,7 +558,10 @@ mod tests { run_in_runtime(async { let mut cooldowns = OriginCooldowns::new(); let outcome = attempt_upload(&mut cooldowns, mock_upload(url, "redirect-me")).await; - assert!(matches!(outcome, AttemptOutcome::Succeeded)); + assert!(matches!( + outcome, + AttemptOutcome::Attempted { succeeded: true } + )); }); redirect_mock.assert(); diff --git a/crates/trident/src/main.rs b/crates/trident/src/main.rs index 0b9d91feb..acd92a102 100644 --- a/crates/trident/src/main.rs +++ b/crates/trident/src/main.rs @@ -1,9 +1,9 @@ -use std::{fs, iter, panic, process::ExitCode}; +use std::{fs, iter, panic, process::ExitCode, time::Duration}; use trident::JournaldTracing; use anyhow::{Context, Error}; use clap::Parser; -use log::{error, info, LevelFilter, Log}; +use log::{error, info, warn, LevelFilter, Log}; use osutils::logging::{filter::LogFilter, multilog::MultiLogger}; use trident::{ @@ -11,9 +11,9 @@ use trident::{ cli::{self, Cli, Commands, GetKind, TridentExitCodes}, init::offline, manual_rollback::{self, utils::ManualRollbackRequestKind}, - run_command, run_reboot_command, save_reboot_operation, validation, BackgroundLog, - BackgroundUploader, DataStore, ExitKind, LogForwarder, Logstream, OperationSource, TraceStream, - Trident, TRIDENT_BACKGROUND_LOG_PATH, + run_command, run_reboot_command, save_reboot_operation, validation, AppInsightsSender, + BackgroundLog, BackgroundUploader, DataStore, ExitKind, LogForwarder, Logstream, + OperationSource, TelemetryUploader, TraceStream, Trident, TRIDENT_BACKGROUND_LOG_PATH, }; use trident_api::{ config::HostConfigurationSource, @@ -407,10 +407,82 @@ fn setup_logging( Ok(logstream) } -fn setup_tracing(args: &Cli) -> Result { +/// Whether the Application Insights tracing layer ended up active on this +/// invocation, and why not when it didn't. Computed by [`setup_tracing`] and +/// surfaced via [`TelemetryStatus::log`] once real logging is available, so +/// operators can tell -- from the logs alone, without reading source -- +/// whether telemetry should be expected to actually reach Application +/// Insights, rather than silently assuming it based on the `Telemetry=` +/// setting alone (a bad/unreachable connection string, for example, fails +/// silently otherwise). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TelemetryStatus { + /// Tracing/telemetry setup does not apply to this command at all (the + /// `Commands::Pytest` arm in [`setup_tracing`], gated behind the + /// `pytest-generator` feature) -- not logged. Only ever constructed + /// when that feature is enabled; every other command now gets a real + /// subscriber installed. + #[cfg_attr(not(feature = "pytest-generator"), allow(dead_code))] + NotApplicable, + /// `Telemetry=OptOut` (the default): telemetry was never attempted. + OptedOut, + /// Opted in, but no usable `AZURE_MONITOR_CONNECTION_STRING` was + /// compiled into this binary at build time (missing, empty, or failed + /// to parse). + NoConnectionString, + /// Opted in with a connection string, but the dedicated telemetry + /// background uploader is unavailable (failed to start, or its handle + /// was already closed). + UploaderUnavailable, + /// Opted in, connection string valid, uploader available: telemetry is + /// actively being sent. + Enabled, +} + +impl TelemetryStatus { + /// Log this status through the real logging pipeline. Must only be + /// called after logging has been initialized (`setup_logging`) -- + /// calling it earlier would silently no-op, since the `log` facade + /// drops everything until a logger is registered. + fn log(self) { + match self { + TelemetryStatus::NotApplicable => {} + TelemetryStatus::OptedOut => { + info!( + "Telemetry: disabled (Telemetry=OptOut, the default, in agent configuration)" + ); + } + TelemetryStatus::NoConnectionString => { + info!( + "Telemetry: opted in, but no usable Application Insights connection string \ + was compiled into this binary -- telemetry is a no-op" + ); + } + TelemetryStatus::UploaderUnavailable => { + warn!( + "Telemetry: opted in, but the telemetry background uploader is \ + unavailable -- telemetry is a no-op" + ); + } + TelemetryStatus::Enabled => { + info!("Telemetry: enabled, sending tracing data to Application Insights"); + } + } + } +} + +fn setup_tracing( + args: &Cli, + telemetry_enabled: bool, + // Dedicated to Application Insights telemetry. This is separate from the + // log-forwarding `BackgroundUploader` so slow telemetry cannot delay real + // log uploads. `None` means telemetry is disabled or unavailable. + telemetry_uploader: Option<&TelemetryUploader>, +) -> Result<(TraceStream, TelemetryStatus), Error> { use tracing_subscriber::{filter, layer::SubscriberExt, Layer, Registry}; let tracestream = TraceStream::default(); + let telemetry_status; // Every command reachable from run_trident needs a subscriber // installed here -- not just the servicing ones -- so that ordinary @@ -473,6 +545,34 @@ fn setup_tracing(args: &Cli) -> Result { } } + // Best-effort Application Insights telemetry: enable it only when + // the user opted in and a usable connection string was compiled in. + // Startup never fails here; telemetry just becomes a no-op, and + // `telemetry_status` records why for later logging. + telemetry_status = if !telemetry_enabled { + TelemetryStatus::OptedOut + } else { + // A missing or closed uploader just leaves telemetry as a + // no-op; tracing setup must still succeed. + match telemetry_uploader.and_then(|u| u.get_handle()) { + Some(handle) => match AppInsightsSender::from_connection_string( + trident::AZURE_MONITOR_CONNECTION_STRING, + handle, + tracestream.installation_id_handle(), + tracestream.servicing_id_handle(), + ) { + Some(sender) => { + // Measure spans before TraceSender performs synchronous HTTP I/O. + layers + .insert(0, Box::new(sender.with_filter(filter::LevelFilter::INFO))); + TelemetryStatus::Enabled + } + None => TelemetryStatus::NoConnectionString, + }, + None => TelemetryStatus::UploaderUnavailable, + } + }; + tracing::subscriber::set_global_default(Registry::default().with(layers)) .context("Failed to set global default subscriber")?; } @@ -484,10 +584,48 @@ fn setup_tracing(args: &Cli) -> Result { // falling through to "no subscriber" the way the commands above // used to. #[cfg(feature = "pytest-generator")] - Commands::Pytest => {} + Commands::Pytest => { + telemetry_status = TelemetryStatus::NotApplicable; + } } - Ok(tracestream) + Ok((tracestream, telemetry_status)) +} + +/// How long to wait for the telemetry uploader to drain before abandoning +/// it. Telemetry must not meaningfully delay shutdown. +const TELEMETRY_SHUTDOWN_DEADLINE: Duration = Duration::from_secs(5); + +/// Maximum number of pending telemetry uploads queued at once. +/// This caps memory if the endpoint is merely slow; once full, the oldest +/// not-yet-attempted event is dropped in favor of the newest. +const TELEMETRY_QUEUE_CAPACITY: usize = 1000; + +/// Wraps a `TelemetryUploader` so every `main` exit path shuts it down with +/// a bounded deadline. `TelemetryUploader`'s own `Drop` waits unboundedly, +/// which is fine for real log delivery but not for telemetry. +struct TelemetryUploaderGuard(Option); + +impl TelemetryUploaderGuard { + fn flush(&self) { + if let Some(uploader) = self.0.as_ref() { + if !uploader.flush_with_deadline(TELEMETRY_SHUTDOWN_DEADLINE) { + log::debug!("Telemetry uploader did not flush before the reboot deadline"); + } + } + } + + fn shutdown(&mut self) { + if let Some(uploader) = self.0.take() { + uploader.shutdown_with_deadline(TELEMETRY_SHUTDOWN_DEADLINE); + } + } +} + +impl Drop for TelemetryUploaderGuard { + fn drop(&mut self) { + self.shutdown(); + } } fn main() -> ExitCode { @@ -503,13 +641,46 @@ fn main() -> ExitCode { } }; + // Whether best-effort Application Insights telemetry is enabled. Loaded + // early (before logging/tracing is set up) since the decision feeds + // directly into setup_tracing(). + let telemetry_enabled = match AgentConfig::load() { + Ok(config) => config.telemetry_enabled(), + Err(error) => { + eprintln!("Failed to load agent configuration, disabling telemetry: {error:?}"); + false + } + }; + + // Application Insights telemetry gets its own dedicated uploader/queue, + // entirely separate from `bg_uploader` (which carries real log + // forwarding): a `TelemetryUploader` rather than a `BackgroundUploader`, + // so a slow-but-successful telemetry endpoint can never build a backlog + // that delays operational log uploads. Failure to start is not fatal: + // telemetry simply becomes a no-op, mirroring failure handling on the + // handle itself. + let telemetry_uploader = telemetry_enabled + .then(|| match TelemetryUploader::new(TELEMETRY_QUEUE_CAPACITY) { + Ok(uploader) => Some(uploader), + Err(e) => { + eprintln!("Failed to initialize telemetry uploader, disabling telemetry: {e:?}"); + None + } + }) + .flatten(); + // Wrapped immediately so every return path in main() below shuts it + // down with a bounded deadline, not TelemetryUploader's own unbounded + // Drop. + let telemetry_uploader = TelemetryUploaderGuard(telemetry_uploader); + // Initialize the telemetry flow - let tracestream = setup_tracing(&args); - if let Err(e) = tracestream { + let tracing_setup = setup_tracing(&args, telemetry_enabled, telemetry_uploader.0.as_ref()); + if let Err(e) = tracing_setup { // Defer to stderr since logging is not yet initialized. eprintln!("Failed to initialize tracing: {e:?}"); return TridentExitCodes::SetupFailed.into(); } + let (tracestream, telemetry_status) = tracing_setup.unwrap(); if let Commands::Daemon { inactivity_timeout, @@ -535,13 +706,15 @@ fn main() -> ExitCode { // Log version on startup info!("Trident version: {}", trident::TRIDENT_VERSION); + telemetry_status.log(); trident::server_main( log_forwarder, *inactivity_timeout, socket_path, logstream.unwrap(), - tracestream.unwrap(), + tracestream, + || telemetry_uploader.flush(), ) } else if let Commands::GrpcClient(client_args) = &args.command { let logstream = setup_logging(&args, &bg_uploader, iter::empty()); @@ -554,8 +727,10 @@ fn main() -> ExitCode { error!("Failed to initialize logstream from environment: {e:?}"); } + telemetry_status.log(); + // Run the client command - trident::client_main(client_args) + trident::client_main(client_args, || telemetry_uploader.flush()) } else { // Initialize the loggers let logstream = setup_logging(&args, &bg_uploader, iter::empty()); @@ -564,15 +739,32 @@ fn main() -> ExitCode { return TridentExitCodes::SetupFailed.into(); } + telemetry_status.log(); + // Invoke Trident - match run_trident(logstream.unwrap(), tracestream.unwrap(), &args) { + match run_trident(logstream.unwrap(), tracestream, &args) { Ok(ExitKind::Done) => {} Err(e) => { error!("{e:?}"); return TridentExitCodes::Failed.into(); } Ok(ExitKind::NeedsReboot) => { - if let Err(e) = run_reboot_command(trident::request_reboot_with_wait) { + // Reuse the just-completed install/update/etc.'s own + // operation_id/command (captured via save_reboot_operation + // just before that command's own run_command scope ended) + // rather than leaving `trident_system_reboot` untagged, or + // minting an unrelated fresh "reboot" identity: the reboot + // is a direct continuation of that same servicing + // operation, not an independent one, so telemetry should + // correlate it back to the same operation_id. + // run_reboot_command also still fires command_error on a + // failed reboot -- falling back to a fresh, untagged + // command if nothing was captured -- matching every other + // command's error-reporting contract instead of silently + // dropping this one on the floor. + if let Err(e) = run_reboot_command(|| { + trident::request_reboot_with_wait_and_flush(|| telemetry_uploader.flush()) + }) { error!("Failed to reboot: {e:?}"); return TridentExitCodes::RebootUnsuccessful.into(); } @@ -582,3 +774,47 @@ fn main() -> ExitCode { TridentExitCodes::Success.into() } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_pre_reboot_flush_keeps_uploader_open_for_errors() { + let mut server = mockito::Server::new(); + let reboot = server + .mock("POST", "/reboot") + .with_status(200) + .expect(1) + .create(); + let error = server + .mock("POST", "/error") + .with_status(200) + .expect(1) + .create(); + let uploader = TelemetryUploader::new(4).unwrap(); + let handle = uploader.get_handle().unwrap(); + let guard = TelemetryUploaderGuard(Some(uploader)); + let url = reqwest::Url::parse(&server.url()).unwrap(); + handle + .upload( + &url.join("reboot").unwrap(), + b"reboot".to_vec(), + Duration::from_secs(2), + None, + ) + .unwrap(); + guard.flush(); + reboot.assert(); + handle + .upload( + &url.join("error").unwrap(), + b"failure".to_vec(), + Duration::from_secs(2), + None, + ) + .unwrap(); + drop(guard); + error.assert(); + } +} diff --git a/crates/trident/src/monitor_metrics.rs b/crates/trident/src/monitor_metrics.rs index 846a5245c..58764bd90 100644 --- a/crates/trident/src/monitor_metrics.rs +++ b/crates/trident/src/monitor_metrics.rs @@ -200,6 +200,11 @@ impl MonitorMetrics { /// It ensures that the thread is stopped gracefully. pub fn stop(&mut self) -> Result<(), Error> { self.stop.store(true, std::sync::atomic::Ordering::SeqCst); + if let Some(handle) = self.join_handle.take() { + handle + .join() + .map_err(|_| Error::msg("Metrics monitoring thread panicked"))?; + } Ok(()) } @@ -506,6 +511,65 @@ mod stat_tests { ); } + #[test] + fn test_stop_waits_for_monitor_completion() { + let stop = Arc::new(AtomicBool::new(false)); + let completed = Arc::new(AtomicBool::new(false)); + let thread_stop = stop.clone(); + let thread_completed = completed.clone(); + let mut monitor = MonitorMetrics { + stop, + metric_count: Arc::new(AtomicU64::new(0)), + join_handle: Some(thread::spawn(move || { + while !thread_stop.load(std::sync::atomic::Ordering::SeqCst) { + thread::yield_now(); + } + thread_completed.store(true, std::sync::atomic::Ordering::SeqCst); + })), + }; + + monitor.stop().unwrap(); + assert!(completed.load(std::sync::atomic::Ordering::SeqCst)); + assert!(monitor.join_handle.is_none()); + monitor.stop().unwrap(); + } + + #[test] + fn test_drop_waits_for_monitor_on_error() { + fn run(completed: Arc) -> Result<(), Error> { + let stop = Arc::new(AtomicBool::new(false)); + let thread_stop = stop.clone(); + let _monitor = MonitorMetrics { + stop, + metric_count: Arc::new(AtomicU64::new(0)), + join_handle: Some(thread::spawn(move || { + while !thread_stop.load(std::sync::atomic::Ordering::SeqCst) { + thread::yield_now(); + } + completed.store(true, std::sync::atomic::Ordering::SeqCst); + })), + }; + Err(Error::msg("servicing failed")) + } + + let completed = Arc::new(AtomicBool::new(false)); + assert!(run(completed.clone()).is_err()); + assert!(completed.load(std::sync::atomic::Ordering::SeqCst)); + } + + #[test] + fn test_stop_reports_monitor_panic() { + let mut monitor = MonitorMetrics { + stop: Arc::new(AtomicBool::new(false)), + metric_count: Arc::new(AtomicU64::new(0)), + join_handle: Some(thread::spawn(|| panic!("monitor failed"))), + }; + + let error = monitor.stop().unwrap_err(); + assert_eq!(error.to_string(), "Metrics monitoring thread panicked"); + monitor.stop().unwrap(); + } + #[test] fn test_join() { let mut monitor = MonitorMetrics { diff --git a/crates/trident/src/reboot.rs b/crates/trident/src/reboot.rs index e8f76474a..73caed62f 100644 --- a/crates/trident/src/reboot.rs +++ b/crates/trident/src/reboot.rs @@ -11,7 +11,14 @@ pub const REBOOT_WAIT_DURATION_SECS: u64 = 600; // 10 minutes /// Calls `request_reboot` and waits for a fixed duration to allow the reboot to /// occur. If the system is still running after the wait, an error is returned. pub fn request_reboot_with_wait() -> Result<(), TridentError> { - request_reboot().structured(ServicingError::Reboot)?; + request_reboot_with_wait_and_flush(|| {}) +} + +/// Drains telemetry after the reboot metric, before systemd can terminate the process. +pub fn request_reboot_with_wait_and_flush( + before_reboot: impl FnOnce(), +) -> Result<(), TridentError> { + request_reboot_with_flush(before_reboot).structured(ServicingError::Reboot)?; thread::sleep(Duration::from_secs(REBOOT_WAIT_DURATION_SECS)); error!( @@ -21,14 +28,33 @@ pub fn request_reboot_with_wait() -> Result<(), TridentError> { } /// Issues a reboot request to systemd after a filesystem sync. -pub fn request_reboot() -> Result<(), Box> { +pub fn request_reboot(before_reboot: impl FnOnce()) -> Result<(), Box> { + request_reboot_with_flush(before_reboot) +} + +fn request_reboot_with_flush(before_reboot: impl FnOnce()) -> Result<(), Box> { // Sync all writes to the filesystem. info!("Syncing filesystem"); nix::unistd::sync(); + request_reboot_inner(before_reboot, || { + Dependency::Systemctl + .cmd() + .env("SYSTEMD_IGNORE_CHROOT", "true") + .arg("reboot") + .run_and_check() + .map(|_| ()) + }) +} + +fn request_reboot_inner( + before_reboot: impl FnOnce(), + request: impl FnOnce() -> Result<(), Box>, +) -> Result<(), Box> { // This trace event will be used with the trident_start event to track the // total time taken for the reboot tracing::info!(metric_name = events::TRIDENT_SYSTEM_REBOOT); + before_reboot(); // IMPORTANT: This message is used by E2E A/B update tests to validate that // a reboot was requested. Do not change or remove without updating the @@ -36,13 +62,64 @@ pub fn request_reboot() -> Result<(), Box> { // tools/storm/utils/trident/trident.go:REBOOTING_LOG_MESSAGE debug!("Requesting reboot"); - Dependency::Systemctl - .cmd() - .env("SYSTEMD_IGNORE_CHROOT", "true") - .arg("reboot") - .run_and_check()?; + request()?; info!("Rebooting system"); Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + use std::{ + cell::Cell, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + }; + + use tracing::{field::Visit, Event, Subscriber}; + use tracing_subscriber::{layer::Context, prelude::*, Layer, Registry}; + + struct RebootMetric(Arc); + + impl Layer for RebootMetric { + fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) { + struct Visitor<'a>(&'a AtomicBool); + impl Visit for Visitor<'_> { + fn record_str(&mut self, field: &tracing::field::Field, value: &str) { + if field.name() == "metric_name" && value == "trident_system_reboot" { + self.0.store(true, Ordering::SeqCst); + } + } + + fn record_debug(&mut self, _: &tracing::field::Field, _: &dyn std::fmt::Debug) {} + } + event.record(&mut Visitor(&self.0)); + } + } + + #[test] + fn test_reboot_metric_is_flushed_before_request() { + let metric_seen = Arc::new(AtomicBool::new(false)); + let flushed = Cell::new(false); + let subscriber = Registry::default().with(RebootMetric(metric_seen.clone())); + tracing::subscriber::with_default(subscriber, || { + request_reboot_inner( + || { + assert!(metric_seen.load(Ordering::SeqCst)); + flushed.set(true); + }, + || { + assert!(flushed.get()); + Ok(()) + }, + ) + .unwrap(); + }); + assert!(flushed.get()); + } +} diff --git a/crates/trident/src/server/mod.rs b/crates/trident/src/server/mod.rs index d24c7279c..a272ac945 100644 --- a/crates/trident/src/server/mod.rs +++ b/crates/trident/src/server/mod.rs @@ -86,6 +86,7 @@ pub fn server_main( default_socket_path: impl AsRef, logstream: Logstream, tracestream: TraceStream, + flush_telemetry: impl FnOnce(), ) -> ExitCode { // Start the Tokio runtime let Ok(runtime) = Builder::new_multi_thread().enable_all().build() else { @@ -168,11 +169,11 @@ pub fn server_main( ExitKind::Done => TridentExitCodes::Success.into(), // Reboot requested - ExitKind::NeedsReboot => reboot(shutdown_signals), + ExitKind::NeedsReboot => reboot(shutdown_signals, flush_telemetry), } } -fn reboot(signals: ShutdownSignals) -> ExitCode { +fn reboot(signals: ShutdownSignals, flush_telemetry: impl FnOnce()) -> ExitCode { // Reuse the just-finished servicing operation's own operation_id/ // command (captured via `save_reboot_operation` from inside // `servicing_request`'s closure, before that operation's own @@ -190,7 +191,7 @@ fn reboot(signals: ShutdownSignals) -> ExitCode { // shutdown signal never arriving) distinguishable in telemetry from a // successful one. let reboot_result: Result<(), TridentError> = operation_context::run_reboot_command(|| { - reboot::request_reboot().structured(ServicingError::Reboot)?; + reboot::request_reboot(flush_telemetry).structured(ServicingError::Reboot)?; // Wait for either a shutdown signal or the reboot timeout. signals diff --git a/docs/Reference/Agent-Configuration.md b/docs/Reference/Agent-Configuration.md index afa12bfda..cc58d28ba 100644 --- a/docs/Reference/Agent-Configuration.md +++ b/docs/Reference/Agent-Configuration.md @@ -24,14 +24,15 @@ DatastorePath=/special/path/to/my-datastore.sqlite ## Telemetry -The `Telemetry` setting records the preference for sending data to -Application Insights. Recording this preference alone does not upload -data: delivery also requires Application Insights runtime integration -and a connection string compiled into the binary using -`AZURE_MONITOR_CONNECTION_STRING`. - -The preference defaults to **OptOut**. To record an opt-in preference, -add a line to the Agent Configuration file: +Trident can optionally send telemetry data to Azure Monitor / Application +Insights. See [Telemetry](./Telemetry.md) for what data is collected and +how delivery failures are handled. This requires an Application Insights +connection string to have been compiled into the Trident binary at build +time (via the `AZURE_MONITOR_CONNECTION_STRING` environment variable); if +no connection string was compiled in, this setting has no effect. + +Telemetry defaults to **disabled** (`OptOut`). To enable it, add a line to +the Agent Configuration file: ``` conf Telemetry=OptIn diff --git a/docs/Reference/Telemetry.md b/docs/Reference/Telemetry.md index 45c0a7d28..313e8e385 100644 --- a/docs/Reference/Telemetry.md +++ b/docs/Reference/Telemetry.md @@ -35,10 +35,23 @@ Servicing paths later copy the current file beside the configured datastore as `trident-metrics--.jsonl` (under `/var/lib/trident` when using the default datastore path). +On top of these local copies, Trident can optionally send this same +best-effort stream of tracing data to Azure Monitor / Application +Insights. See [Agent Configuration](./Agent-Configuration.md) for how to +enable it. + +Telemetry defaults to **disabled** (`OptOut`). To enable it, add a line to +the Agent Configuration file: + +``` conf +Telemetry=OptIn +``` + ## Host Metadata -Every recorded event also includes as much of the following host metadata -as is available at the time: +Every event sent also includes as much of the following host metadata as +is available at the time, so operators should be aware this leaves the +host along with the metrics/spans themselves: - `os_release`: the `VERSION` field from `/etc/os-release`. - `kernel_version`: the running kernel release (`uname -r`). @@ -185,3 +198,13 @@ Admission-control rejections for a busy daemon (connection or servicing lock contention) or an unprovisioned host intentionally emit no `command_error` from either the daemon or the client. Consequently, `command_error` is not a count of every failed invocation. + +## Delivery + +Telemetry delivery is always best-effort and never affects servicing +outcomes, but failures are not all logged at the same level: a failure to +serialize an event, or to enqueue it because the background uploader has +already shut down, is logged at trace level, while a failure to actually +deliver an event (e.g. no network connectivity, or a non-2xx response from +Application Insights) is logged at error level, so operators can find +remote-delivery problems in normal logs.