Add shared OTLP telemetry and propagate command traces across agents

wakey-core/Cargo.toml

@@ -5,11 +5,28 @@ edition = "2024"
 publish = ["gitea"]
 
 [dependencies]
+anyhow = { workspace = true, optional = true }
+opentelemetry = { version = "0.31", optional = true }
+opentelemetry-otlp = { version = "0.31", default-features = false, features = ["grpc-tonic", "trace", "logs", "metrics"], optional = true }
+opentelemetry_sdk = { version = "0.31", optional = true }
+opentelemetry-appender-tracing = { version = "0.31", features = ["experimental_use_tracing_span_context"], optional = true }
+tracing = { workspace = true, optional = true }
+tracing-opentelemetry = { version = "0.32", optional = true }
+tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt", "json", "time"], optional = true }
+time = { version = "0.3", features = ["formatting", "local-offset"], optional = true }
+tonic = { version = "0.14", default-features = false, optional = true }
 macaddr = { workspace = true }
 serde = { workspace = true }
 serde_with = { version = "3", features = ["json"] }
 strum = { version = "0", features = ["derive", "strum_macros"] }
 thiserror = "2"
 
+[features]
+telemetry = ["dep:anyhow", "dep:opentelemetry", "dep:opentelemetry-otlp", "dep:opentelemetry_sdk", "dep:opentelemetry-appender-tracing", "dep:tracing", "dep:tracing-opentelemetry", "dep:tracing-subscriber", "dep:time", "dep:tonic"]
+
 [dev-dependencies]
 serde_json.workspace = true
+tokio = { version = "1", features = ["macros", "rt-multi-thread", "net", "sync", "time"] }
+tokio-stream = { version = "0.1", features = ["net"] }
+tonic = "0.14"
+opentelemetry-proto = { version = "0.31", default-features = false, features = ["gen-tonic", "trace", "logs", "metrics"] }

wakey-core/src/lib.rs

@@ -3,4 +3,7 @@
 pub mod model;
 pub mod parse;
 
+#[cfg(feature = "telemetry")]
+pub mod telemetry;
+
 pub use model::*;

wakey-core/src/telemetry.rs

@@ -0,0 +1,409 @@
+//! Shared, opt-in OTLP setup. Local logging remains available without a collector.
+
+use std::collections::BTreeMap;
+use std::sync::OnceLock;
+use std::time::{Duration, Instant};
+
+use anyhow::{Context, Result};
+use opentelemetry::{
+    KeyValue, global,
+    trace::{TraceContextExt as _, TracerProvider as _},
+};
+use opentelemetry_otlp::WithExportConfig;
+use opentelemetry_otlp::WithTonicConfig;
+use opentelemetry_sdk::{
+    Resource, logs::SdkLoggerProvider, metrics::SdkMeterProvider, trace::SdkTracerProvider,
+};
+use serde::{Deserialize, Serialize};
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+use tracing_subscriber::{EnvFilter, Layer, fmt, layer::SubscriberExt, util::SubscriberInitExt};
+
+/// A single gRPC receiver accepts all three signals. No endpoint means local-only.
+#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(default)]
+pub struct TelemetryConfig {
+    pub otlp_endpoint: Option<String>,
+    /// Name of the environment variable holding the complete Authorization header.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub authorization_env: Option<String>,
+    pub service_name: String,
+    pub json_logs: bool,
+}
+
+/// Owns providers until all instrumented work has stopped; flush off the async workers.
+#[derive(Default)]
+pub struct TelemetryGuard {
+    traces: Option<SdkTracerProvider>,
+    logs: Option<SdkLoggerProvider>,
+    metrics: Option<SdkMeterProvider>,
+}
+
+impl TelemetryGuard {
+    pub fn shutdown(self) {
+        // Export failures must not turn an otherwise successful daemon exit into a failure.
+        if let Some(provider) = self.metrics {
+            let _ = provider.shutdown();
+        }
+        if let Some(provider) = self.logs {
+            let _ = provider.shutdown();
+        }
+        if let Some(provider) = self.traces {
+            let _ = provider.shutdown();
+        }
+    }
+}
+
+pub fn init(
+    config: &TelemetryConfig,
+    service: &str,
+    default_filter: &str,
+    instance: Option<&str>,
+) -> Result<TelemetryGuard> {
+    let filter =
+        EnvFilter::try_from_default_env().or_else(|_| EnvFilter::try_new(default_filter))?;
+    let service = if config.service_name.is_empty() {
+        service
+    } else {
+        &config.service_name
+    };
+    let mut guard = TelemetryGuard::default();
+    if let Some(endpoint) = config
+        .otlp_endpoint
+        .as_deref()
+        .filter(|s| !s.trim().is_empty())
+    {
+        let metadata = authorization_metadata(config.authorization_env.as_deref(), |name| {
+            std::env::var(name).ok()
+        })?;
+        let mut resource = Resource::builder_empty().with_service_name(service.to_owned());
+        if let Some(instance) = instance {
+            resource =
+                resource.with_attribute(KeyValue::new("service.instance.id", instance.to_owned()));
+        }
+        let resource = resource.build();
+        // Small fixed queues are intentional: a disconnected collector cannot grow router RAM.
+        let trace_exporter = opentelemetry_otlp::SpanExporter::builder()
+            .with_tonic()
+            .with_metadata(metadata.clone())
+            .with_endpoint(endpoint)
+            .with_timeout(Duration::from_secs(3))
+            .build()
+            .context("OTLP trace exporter")?;
+        let trace_processor = opentelemetry_sdk::trace::BatchSpanProcessor::builder(trace_exporter)
+            .with_batch_config(
+                opentelemetry_sdk::trace::BatchConfigBuilder::default()
+                    .with_max_queue_size(256)
+                    .with_max_export_batch_size(64)
+                    .build(),
+            )
+            .build();
+        let traces = SdkTracerProvider::builder()
+            .with_span_processor(trace_processor)
+            .with_resource(resource.clone())
+            .build();
+        let log_exporter = opentelemetry_otlp::LogExporter::builder()
+            .with_tonic()
+            .with_metadata(metadata.clone())
+            .with_endpoint(endpoint)
+            .with_timeout(Duration::from_secs(3))
+            .build()
+            .context("OTLP log exporter")?;
+        let log_processor = opentelemetry_sdk::logs::BatchLogProcessor::builder(log_exporter)
+            .with_batch_config(
+                opentelemetry_sdk::logs::BatchConfigBuilder::default()
+                    .with_max_queue_size(256)
+                    .with_max_export_batch_size(64)
+                    .build(),
+            )
+            .build();
+        let logs = SdkLoggerProvider::builder()
+            .with_log_processor(log_processor)
+            .with_resource(resource.clone())
+            .build();
+        let metric_exporter = opentelemetry_otlp::MetricExporter::builder()
+            .with_tonic()
+            .with_metadata(metadata)
+            .with_endpoint(endpoint)
+            .with_timeout(Duration::from_secs(3))
+            .build()
+            .context("OTLP metric exporter")?;
+        let reader = opentelemetry_sdk::metrics::PeriodicReader::builder(metric_exporter)
+            .with_interval(Duration::from_secs(30))
+            .build();
+        let metrics = SdkMeterProvider::builder()
+            .with_reader(reader)
+            .with_resource(resource)
+            .build();
+        global::set_tracer_provider(traces.clone());
+        global::set_meter_provider(metrics.clone());
+        let _ = METRICS.set(Instruments::new());
+        let otel = tracing_opentelemetry::layer().with_tracer(traces.tracer(service.to_owned()));
+        // Exporter failures must remain local, never feed back into the failing exporter.
+        let logs_layer =
+            opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(&logs)
+                .with_filter(tracing_subscriber::filter::filter_fn(|metadata| {
+                    !metadata.target().starts_with("opentelemetry")
+                }));
+        let subscriber = tracing_subscriber::registry()
+            .with(filter)
+            .with(otel)
+            .with(logs_layer);
+        if config.json_logs {
+            subscriber
+                .with(fmt::layer().json().with_timer(local_timer()))
+                .try_init()?;
+        } else {
+            subscriber
+                .with(fmt::layer().with_timer(local_timer()))
+                .try_init()?;
+        }
+        guard = TelemetryGuard {
+            traces: Some(traces),
+            logs: Some(logs),
+            metrics: Some(metrics),
+        };
+    } else if config.json_logs {
+        tracing_subscriber::registry()
+            .with(filter)
+            .with(fmt::layer().json().with_timer(local_timer()))
+            .try_init()?;
+    } else {
+        tracing_subscriber::registry()
+            .with(filter)
+            .with(fmt::layer().with_timer(local_timer()))
+            .try_init()?;
+    }
+    Ok(guard)
+}
+
+fn authorization_metadata(
+    name: Option<&str>,
+    lookup: impl FnOnce(&str) -> Option<String>,
+) -> Result<tonic::metadata::MetadataMap> {
+    let mut metadata = tonic::metadata::MetadataMap::new();
+    let Some(name) = name else {
+        return Ok(metadata);
+    };
+    anyhow::ensure!(
+        !name.is_empty()
+            && name.bytes().all(|c| c.is_ascii_alphanumeric() || c == b'_')
+            && !name.as_bytes()[0].is_ascii_digit(),
+        "telemetry authorization_env must be a valid environment variable name"
+    );
+    let raw = lookup(name).with_context(|| {
+        format!("telemetry authorization environment variable {name} is missing or not UTF-8")
+    })?;
+    let header = raw.trim();
+    anyhow::ensure!(
+        !header.is_empty(),
+        "telemetry authorization environment variable {name} is empty"
+    );
+    let valid_scheme = header.split_once(' ').is_some_and(|(scheme, credentials)| {
+        (scheme.eq_ignore_ascii_case("basic") || scheme.eq_ignore_ascii_case("bearer"))
+            && !credentials.trim().is_empty()
+    });
+    anyhow::ensure!(
+        valid_scheme && !raw.contains(['\r', '\n']),
+        "telemetry authorization environment variable {name} must contain a single Basic or Bearer header"
+    );
+    // Never include the supplied value/parser error in diagnostics. Sensitive
+    // metadata also hides it if a dependency Debug-prints an outgoing request.
+    let mut value = tonic::metadata::MetadataValue::try_from(header).map_err(|_| {
+        anyhow::anyhow!("invalid authorization metadata in environment variable {name}")
+    })?;
+    value.set_sensitive(true);
+    metadata.insert("authorization", value);
+    Ok(metadata)
+}
+
+fn local_timer() -> fmt::time::OffsetTime<time::format_description::well_known::Rfc3339> {
+    fmt::time::OffsetTime::new(
+        time::UtcOffset::current_local_offset().unwrap_or(time::UtcOffset::UTC),
+        time::format_description::well_known::Rfc3339,
+    )
+}
+
+// Construct instruments once, after the real provider exists. Disabled telemetry
+// therefore does no per-PTY-read allocations or instrument registration.
+static METRICS: OnceLock<Instruments> = OnceLock::new();
+struct Instruments {
+    operations: opentelemetry::metrics::Counter<u64>,
+    duration: opentelemetry::metrics::Histogram<f64>,
+    snapshot_devices: opentelemetry::metrics::Counter<u64>,
+    terminal_active: opentelemetry::metrics::UpDownCounter<i64>,
+    terminal_input: opentelemetry::metrics::Counter<u64>,
+    terminal_output: opentelemetry::metrics::Counter<u64>,
+    terminal_backpressure: opentelemetry::metrics::Counter<u64>,
+}
+impl Instruments {
+    fn new() -> Self {
+        let meter = global::meter("wakey");
+        Self {
+            operations: meter.u64_counter("wakey.operations").build(),
+            duration: meter
+                .f64_histogram("wakey.operation.duration")
+                .with_unit("s")
+                .build(),
+            snapshot_devices: meter.u64_counter("wakey.snapshot.devices").build(),
+            terminal_active: meter.i64_up_down_counter("wakey.terminal.active").build(),
+            terminal_input: meter
+                .u64_counter("wakey.terminal.input.bytes")
+                .with_unit("By")
+                .build(),
+            terminal_output: meter
+                .u64_counter("wakey.terminal.output.bytes")
+                .with_unit("By")
+                .build(),
+            terminal_backpressure: meter.u64_counter("wakey.terminal.backpressure").build(),
+        }
+    }
+}
+
+/// W3C context only, never arbitrary application fields or credentials.
+pub fn current_trace_context() -> BTreeMap<String, String> {
+    use opentelemetry::propagation::TextMapPropagator;
+    let mut carrier = std::collections::HashMap::new();
+    opentelemetry_sdk::propagation::TraceContextPropagator::new()
+        .inject_context(&tracing::Span::current().context(), &mut carrier);
+    carrier.into_iter().collect()
+}
+
+pub fn set_remote_parent(span: &tracing::Span, carrier: &BTreeMap<String, String>) {
+    use opentelemetry::propagation::TextMapPropagator;
+    let carrier: std::collections::HashMap<String, String> = carrier
+        .iter()
+        .filter(|(key, value)| {
+            matches!(key.as_str(), "traceparent" | "tracestate") && value.len() <= 512
+        })
+        .map(|(k, v)| (k.clone(), v.clone()))
+        .collect();
+    let context = opentelemetry_sdk::propagation::TraceContextPropagator::new().extract(&carrier);
+    if context.has_active_span() {
+        let _ = span.set_parent(context);
+    }
+}
+
+/// Records completion on every return path, without high-cardinality identifiers.
+pub struct Operation {
+    started: Instant,
+    name: &'static str,
+    outcome: &'static str,
+}
+
+impl Operation {
+    pub fn new(name: &'static str) -> Self {
+        Self {
+            started: Instant::now(),
+            name,
+            outcome: "error",
+        }
+    }
+    pub fn outcome(&mut self, outcome: &'static str) {
+        self.outcome = outcome;
+    }
+}
+
+impl Drop for Operation {
+    fn drop(&mut self) {
+        let Some(metrics) = METRICS.get() else {
+            return;
+        };
+        let labels = [
+            KeyValue::new("operation", self.name),
+            KeyValue::new("outcome", self.outcome),
+        ];
+        metrics.operations.add(1, &labels);
+        metrics
+            .duration
+            .record(self.started.elapsed().as_secs_f64(), &labels);
+    }
+}
+
+pub enum Counter {
+    SnapshotDevices,
+    TerminalInputBytes,
+    TerminalOutputBytes,
+    TerminalBackpressure,
+}
+
+pub fn count(counter: Counter, value: u64) {
+    let Some(metrics) = METRICS.get() else {
+        return;
+    };
+    let counter = match counter {
+        Counter::SnapshotDevices => &metrics.snapshot_devices,
+        Counter::TerminalInputBytes => &metrics.terminal_input,
+        Counter::TerminalOutputBytes => &metrics.terminal_output,
+        Counter::TerminalBackpressure => &metrics.terminal_backpressure,
+    };
+    counter.add(value, &[]);
+}
+
+/// An active session contributes until its worker finishes or is cancelled.
+pub struct ActiveTerminal;
+impl ActiveTerminal {
+    pub fn new() -> Self {
+        if let Some(metrics) = METRICS.get() {
+            metrics.terminal_active.add(1, &[]);
+        }
+        Self
+    }
+}
+impl Default for ActiveTerminal {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+impl Drop for ActiveTerminal {
+    fn drop(&mut self) {
+        if let Some(metrics) = METRICS.get() {
+            metrics.terminal_active.add(-1, &[]);
+        }
+    }
+}
+
+#[cfg(test)]
+mod auth_tests {
+    use super::*;
+
+    #[test]
+    fn authorization_supports_both_schemes_and_redacts_debug() {
+        for header in ["Basic dGVzdDp0ZXN0", "Bearer test-token"] {
+            let metadata =
+                authorization_metadata(Some("OTLP_AUTH"), |_| Some(header.into())).unwrap();
+            assert_eq!(
+                metadata.get("authorization").unwrap().to_str().unwrap(),
+                header
+            );
+            assert!(!format!("{metadata:?}").contains(header));
+        }
+        assert!(
+            authorization_metadata(None, |_| panic!("no credential lookup"))
+                .unwrap()
+                .is_empty()
+        );
+    }
+
+    #[test]
+    fn missing_empty_and_invalid_auth_fail_without_echoing_secrets() {
+        for value in [
+            None,
+            Some(""),
+            Some("   "),
+            Some("Basic "),
+            Some("Digest private-secret"),
+            Some("Basic private-secret\r\n"),
+            Some("Bearer private-secret\u{7f}"),
+        ] {
+            let error = authorization_metadata(Some("OTLP_AUTH"), |_| value.map(str::to_owned))
+                .unwrap_err();
+            assert!(!format!("{error:#}").contains("private-secret"));
+        }
+        assert!(
+            authorization_metadata(Some("BAD=NAME"), |_| panic!(
+                "invalid name must not be looked up"
+            ))
+            .is_err()
+        );
+    }
+}

wakey-core/tests/telemetry.rs

@@ -0,0 +1,196 @@
+#![cfg(feature = "telemetry")]
+
+use opentelemetry_proto::tonic::collector::{
+    logs::v1 as logs, metrics::v1 as metrics, trace::v1 as traces,
+};
+use std::sync::{Arc, Mutex};
+use tonic::{Request, Response, Status};
+use wakey_core::telemetry::{self, TelemetryConfig};
+
+#[derive(Clone, Default)]
+struct Collector {
+    traces: Arc<Mutex<Vec<traces::ExportTraceServiceRequest>>>,
+    logs: Arc<Mutex<Vec<logs::ExportLogsServiceRequest>>>,
+    metrics: Arc<Mutex<Vec<metrics::ExportMetricsServiceRequest>>>,
+}
+
+fn check_auth<T>(request: &Request<T>) -> Result<(), Status> {
+    if request
+        .metadata()
+        .get("authorization")
+        .and_then(|v| v.to_str().ok())
+        != Some("Basic dGVzdDp0ZXN0")
+    {
+        return Err(Status::unauthenticated(
+            "test receiver requires authorization",
+        ));
+    }
+    Ok(())
+}
+
+#[tonic::async_trait]
+impl traces::trace_service_server::TraceService for Collector {
+    async fn export(
+        &self,
+        request: Request<traces::ExportTraceServiceRequest>,
+    ) -> Result<Response<traces::ExportTraceServiceResponse>, Status> {
+        check_auth(&request)?;
+        self.traces.lock().unwrap().push(request.into_inner());
+        Ok(Response::new(Default::default()))
+    }
+}
+#[tonic::async_trait]
+impl logs::logs_service_server::LogsService for Collector {
+    async fn export(
+        &self,
+        request: Request<logs::ExportLogsServiceRequest>,
+    ) -> Result<Response<logs::ExportLogsServiceResponse>, Status> {
+        check_auth(&request)?;
+        self.logs.lock().unwrap().push(request.into_inner());
+        Ok(Response::new(Default::default()))
+    }
+}
+#[tonic::async_trait]
+impl metrics::metrics_service_server::MetricsService for Collector {
+    async fn export(
+        &self,
+        request: Request<metrics::ExportMetricsServiceRequest>,
+    ) -> Result<Response<metrics::ExportMetricsServiceResponse>, Status> {
+        check_auth(&request)?;
+        self.metrics.lock().unwrap().push(request.into_inner());
+        Ok(Response::new(Default::default()))
+    }
+}
+
+/// Shutdown is the synchronization point: no sleeps or export-interval polling.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn exports_all_signals_and_preserves_remote_command_parent() {
+    // Exercise the real global subscriber without mutating this process's environment
+    // or inheriting a developer's RUST_LOG=warn/off setting.
+    if std::env::var_os("WAKEY_OTLP_TEST_CHILD").is_none() {
+        let status = std::process::Command::new(std::env::current_exe().unwrap())
+            .args([
+                "--exact",
+                "exports_all_signals_and_preserves_remote_command_parent",
+                "--nocapture",
+            ])
+            .env("WAKEY_OTLP_TEST_CHILD", "1")
+            .env("RUST_LOG", "info")
+            .env("WAKEY_OTLP_TEST_AUTH", "Basic dGVzdDp0ZXN0")
+            .status()
+            .unwrap();
+        assert!(status.success());
+        return;
+    }
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let endpoint = format!("http://{}", listener.local_addr().unwrap());
+    let collector = Collector::default();
+    let (stop_tx, stop_rx) = tokio::sync::oneshot::channel();
+    let server = tonic::transport::Server::builder()
+        .add_service(traces::trace_service_server::TraceServiceServer::new(
+            collector.clone(),
+        ))
+        .add_service(logs::logs_service_server::LogsServiceServer::new(
+            collector.clone(),
+        ))
+        .add_service(metrics::metrics_service_server::MetricsServiceServer::new(
+            collector.clone(),
+        ))
+        .serve_with_incoming_shutdown(
+            tokio_stream::wrappers::TcpListenerStream::new(listener),
+            async {
+                let _ = stop_rx.await;
+            },
+        );
+    let server = tokio::spawn(server);
+    let guard = telemetry::init(
+        &TelemetryConfig {
+            otlp_endpoint: Some(endpoint),
+            authorization_env: Some("WAKEY_OTLP_TEST_AUTH".into()),
+            service_name: "wakey-test".into(),
+            json_logs: false,
+        },
+        "unused",
+        "info",
+        Some("test-agent"),
+    )
+    .unwrap();
+
+    let parent = tracing::info_span!("relay_command");
+    let context = parent.in_scope(telemetry::current_trace_context);
+    assert!(context.contains_key("traceparent"));
+    let child = tracing::info_span!("execute_command");
+    telemetry::set_remote_parent(&child, &context);
+    child.in_scope(|| {
+        tracing::info!(command = "inventory", "command complete");
+        let mut operation = telemetry::Operation::new("inventory");
+        operation.outcome("ok");
+        telemetry::count(telemetry::Counter::SnapshotDevices, 3);
+        let _active = telemetry::ActiveTerminal::new();
+    });
+    drop(child);
+    drop(parent);
+    tracing::warn!(target: "opentelemetry_sdk", "synthetic exporter failure remains local");
+    let invalid_parent = tracing::info_span!("invalid_remote_parent");
+    telemetry::set_remote_parent(
+        &invalid_parent,
+        &[("traceparent".into(), "malformed".into())].into(),
+    );
+    drop(invalid_parent);
+    tokio::task::spawn_blocking(move || guard.shutdown())
+        .await
+        .unwrap();
+    let _ = stop_tx.send(());
+    server.await.unwrap().unwrap();
+
+    let traces = collector.traces.lock().unwrap();
+    let spans: Vec<_> = traces
+        .iter()
+        .flat_map(|r| &r.resource_spans)
+        .flat_map(|r| &r.scope_spans)
+        .flat_map(|s| &s.spans)
+        .collect();
+    let parent = spans.iter().find(|s| s.name == "relay_command").unwrap();
+    let child = spans.iter().find(|s| s.name == "execute_command").unwrap();
+    assert_eq!(parent.trace_id, child.trace_id);
+    assert_eq!(parent.span_id, child.parent_span_id);
+    let invalid_parent = spans
+        .iter()
+        .find(|s| s.name == "invalid_remote_parent")
+        .unwrap();
+    assert!(invalid_parent.parent_span_id.is_empty());
+    assert_ne!(parent.trace_id, invalid_parent.trace_id);
+    let logs = collector.logs.lock().unwrap();
+    let records: Vec<_> = logs
+        .iter()
+        .flat_map(|r| &r.resource_logs)
+        .flat_map(|r| &r.scope_logs)
+        .flat_map(|s| &s.log_records)
+        .collect();
+    assert!(
+        records
+            .iter()
+            .any(|r| r.trace_id == child.trace_id && r.span_id == child.span_id)
+    );
+    assert!(
+        !records
+            .iter()
+            .any(|r| format!("{:?}", r.body).contains("synthetic exporter failure"))
+    );
+    let metrics = collector.metrics.lock().unwrap();
+    let names: Vec<_> = metrics
+        .iter()
+        .flat_map(|r| &r.resource_metrics)
+        .flat_map(|r| &r.scope_metrics)
+        .flat_map(|s| &s.metrics)
+        .map(|m| m.name.as_str())
+        .collect();
+    for name in [
+        "wakey.operations",
+        "wakey.operation.duration",
+        "wakey.snapshot.devices",
+        "wakey.terminal.active",
+    ] {
+        assert!(names.contains(&name), "missing {name}");
+    }
+}