Add shared OTLP telemetry and propagate command traces across agents

wakey-agent/Cargo.toml

@@ -32,13 +32,7 @@ tokio = { version = "1", features = [
 tokio-tungstenite = { version = "0.29", features = ["rustls-tls-webpki-roots"] }
 toml = "1.1"
 tracing = "0.1"
-tracing-subscriber = { version = "0.3", features = [
-    "env-filter",
-    "fmt",
-    "time",
-] }
-time = { version = "0.3", features = ["formatting", "local-offset"] }
 url = "2"
 vt100 = { package = "wakey-vt100", path = "../vendor/vt100", registry = "gitea", version = "=0.16.2-wakey.2" }
 wakey = { path = "..", registry = "gitea", version = "0" }
-wakey-core = { path = "../wakey-core", registry = "gitea", version = "0" }
+wakey-core = { path = "../wakey-core", registry = "gitea", version = "0", features = ["telemetry"] }

wakey-agent/src/config.rs

@@ -42,6 +42,8 @@ pub struct AgentConfig {
     pub observation_store_path: PathBuf,
     #[serde(default)]
     pub terminal: TerminalConfig,
+    #[serde(default)]
+    pub telemetry: wakey_core::telemetry::TelemetryConfig,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
@@ -88,6 +90,7 @@ pub static DEFAULT_CONFIG: LazyLock<AgentConfig> = LazyLock::new(|| AgentConfig
     mac_name_cache_path: default_mac_name_cache_path(),
     observation_store_path: default_observation_store_path(),
     terminal: TerminalConfig::default(),
+    telemetry: wakey_core::telemetry::TelemetryConfig::default(),
 });
 
 impl fmt::Debug for AgentConfig {
@@ -111,6 +114,7 @@ impl fmt::Debug for AgentConfig {
             .field("mac_name_cache_path", &self.mac_name_cache_path)
             .field("observation_store_path", &self.observation_store_path)
             .field("terminal", &self.terminal)
+            .field("telemetry", &self.telemetry)
             .finish()
     }
 }
@@ -274,6 +278,12 @@ mod tests {
             dhcp_leases_path: "/tmp/test-dhcp.leases".into(),
             mac_name_cache_path: "/tmp/test-names.json".into(),
             observation_store_path: "/tmp/test-observations.json".into(),
+            telemetry: wakey_core::telemetry::TelemetryConfig {
+                otlp_endpoint: Some("http://10.0.0.1:4317".into()),
+                authorization_env: Some("WAKEY_OTLP_AUTHORIZATION".into()),
+                service_name: "router-agent".into(),
+                json_logs: true,
+            },
             terminal: TerminalConfig {
                 enabled: true,
                 shell: "/bin/sh".into(),

wakey-agent/src/enroll.rs

@@ -151,6 +151,12 @@ mod tests {
             dhcp_leases_path: "/tmp/custom-dhcp.leases".into(),
             mac_name_cache_path: "/tmp/custom-names.json".into(),
             observation_store_path: "/tmp/custom-observations.json".into(),
+            telemetry: wakey_core::telemetry::TelemetryConfig {
+                otlp_endpoint: Some("http://10.0.0.1:4317".into()),
+                authorization_env: Some("WAKEY_OTLP_AUTHORIZATION".into()),
+                service_name: "router-agent".into(),
+                json_logs: true,
+            },
             terminal: crate::config::TerminalConfig {
                 enabled: true,
                 shell: "/bin/ash".into(),

wakey-agent/src/main.rs

@@ -15,9 +15,31 @@ use wakey_agent::protocol;
 #[tokio::main]
 async fn main() -> Result<()> {
     let cli = Cli::parse();
-    tracing::init(cli.verbose);
     let global_config = cli.config.as_deref();
 
+    // Only the long-lived daemon exports: hotplug/enroll invocations stay cheap.
+    let daemon_config = if let Command::Serve(args) = &cli.command {
+        Some(config::load_config(global_config.unwrap_or(&args.config))?)
+    } else {
+        None
+    };
+    let local_only = wakey_core::telemetry::TelemetryConfig::default();
+    let telemetry = tracing::init(
+        cli.verbose,
+        daemon_config
+            .as_ref()
+            .map(|c| &c.telemetry)
+            .unwrap_or(&local_only),
+        daemon_config.as_ref().map(|c| c.agent_id.as_str()),
+    )?;
+
+    let result = run(cli).await;
+    tokio::task::spawn_blocking(move || telemetry.shutdown()).await?;
+    result
+}
+
+async fn run(cli: Cli) -> Result<()> {
+    let global_config = cli.config.as_deref();
     match cli.command {
         Command::Serve(mut args) => {
             if let Some(config) = global_config {

wakey-agent/src/protocol.rs

@@ -325,6 +325,8 @@ pub enum ServerMessage {
     Command {
         request_id: RequestId,
         command: AgentCommand,
+        #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
+        trace_context: std::collections::BTreeMap<String, String>,
     },
     SyncDeviceSnapshot,
     SyncTerminalSessions,
@@ -354,6 +356,7 @@ mod tests {
             command: AgentCommand::Leases(LeasesRequest {
                 include_state: true,
             }),
+            trace_context: Default::default(),
         };
 
         let json = serde_json::to_string(&msg).expect("serialize");
@@ -367,6 +370,36 @@ mod tests {
         assert!(err.contains("must not be empty"));
     }
 
+    #[test]
+    fn old_commands_without_trace_context_remain_valid() {
+        let message: ServerMessage = serde_json::from_str(r#"{"type":"command","request_id":"req-old","command":{"kind":"leases","include_state":false}}"#).expect("old command");
+        let ServerMessage::Command { trace_context, .. } = message else {
+            panic!("command expected")
+        };
+        assert!(trace_context.is_empty());
+    }
+
+    #[test]
+    fn command_trace_context_roundtrips() {
+        let message = ServerMessage::Command {
+            request_id: RequestId::try_from("req-traced".to_string()).unwrap(),
+            command: AgentCommand::Leases(LeasesRequest {
+                include_state: false,
+            }),
+            trace_context: [(
+                "traceparent".into(),
+                "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01".into(),
+            )]
+            .into(),
+        };
+        let ServerMessage::Command { trace_context, .. } =
+            serde_json::from_str(&serde_json::to_string(&message).unwrap()).unwrap()
+        else {
+            panic!("command expected")
+        };
+        assert!(trace_context.contains_key("traceparent"));
+    }
+
     #[test]
     fn terminal_id_deserialization_enforces_validation() {
         let valid: TerminalId = serde_json::from_str(r#""term-1""#).expect("valid terminal id");

wakey-agent/src/serve.rs

@@ -29,6 +29,8 @@ pub async fn serve(args: ServeArgs) -> Result<()> {
     {
         use tokio::signal::unix::{SignalKind, signal};
         let mut hup = signal(SignalKind::hangup()).context("failed to install SIGHUP handler")?;
+        let mut terminate =
+            signal(SignalKind::terminate()).context("failed to install SIGTERM handler")?;
 
         loop {
             tokio::select! {
@@ -37,9 +39,17 @@ pub async fn serve(args: ServeArgs) -> Result<()> {
                     worker.abort();
                     break;
                 }
+                _ = terminate.recv() => {
+                    info!("SIGTERM received; shutting down wakey-agent");
+                    worker.abort();
+                    break;
+                }
                 _ = hup.recv() => {
                     match config::load_config(&args.config) {
                         Ok(new_cfg) => {
+                            if new_cfg.telemetry != cfg.telemetry {
+                                warn!("telemetry configuration changed; restart the daemon to apply it");
+                            }
                             cfg = new_cfg;
                             info!(config_path = %args.config.display(), agent_id = %cfg.agent_id, "reload requested; restarting session with updated config");
                             worker.abort();
@@ -67,6 +77,8 @@ pub async fn serve(args: ServeArgs) -> Result<()> {
         worker.abort();
     }
 
+    // Await the cancelled session task before flushing its telemetry.
+    let _ = worker.await;
     let _ = remove_pid_file(&pid_file);
     Ok(())
 }

wakey-agent/src/session.rs

@@ -195,11 +195,13 @@ where
     .await
 }
 
+#[tracing::instrument(name = "collect_device_snapshot", skip_all)]
 async fn send_device_snapshot_ws<S>(sink: &mut S, config: &AgentConfig) -> Result<usize>
 where
     S: SinkExt<Message> + Unpin,
     <S as futures_util::Sink<Message>>::Error: std::error::Error + Send + Sync + 'static,
 {
+    let mut operation = wakey_core::telemetry::Operation::new("snapshot");
     let req = crate::protocol::InventoryRequest {
         query: None,
         name: None,
@@ -221,6 +223,11 @@ where
     )
     .await?;
     debug!(agent_id = %config.agent_id, devices = count, "sent device snapshot over websocket");
+    operation.outcome("ok");
+    wakey_core::telemetry::count(
+        wakey_core::telemetry::Counter::SnapshotDevices,
+        count as u64,
+    );
     Ok(count)
 }
 
@@ -244,11 +251,16 @@ where
         ServerMessage::Command {
             request_id,
             command,
+            trace_context,
         } => {
             let kind = command_kind(&command);
+            let span = info_span!("execute_command", command = kind, request_id = %request_id);
+            wakey_core::telemetry::set_remote_parent(&span, &trace_context);
+            let mut operation = wakey_core::telemetry::Operation::new(kind);
             info!(request_id = %request_id, command = %kind, "received command from control-plane");
-            match dispatch_command(command, config).await {
+            match tracing::Instrument::instrument(dispatch_command(command, config), span).await {
                 Ok(result) => {
+                    operation.outcome("ok");
                     info!(request_id = %request_id, command = %kind, "command execution completed");
                     send_command_result(sink, request_id, kind, result).await?;
                 }

wakey-agent/src/terminal.rs

@@ -80,6 +80,10 @@ impl RelayOutputBuffer {
         match sender.try_send(frame) {
             Ok(()) => {}
             Err(mpsc::error::TrySendError::Full(frame)) => {
+                wakey_core::telemetry::count(
+                    wakey_core::telemetry::Counter::TerminalBackpressure,
+                    1,
+                );
                 self.pending.push_back(frame);
             }
             Err(mpsc::error::TrySendError::Closed(_)) => self.disconnect(),
@@ -169,6 +173,7 @@ impl TerminalManager {
         )
     }
 
+    #[tracing::instrument(name = "terminal_open", skip_all, fields(terminal_id = %terminal_id))]
     pub fn open(
         &self,
         config: &AgentConfig,
@@ -214,6 +219,7 @@ impl TerminalManager {
         let active = Arc::downgrade(&self.active);
         let events = self.events.clone();
         tokio::spawn(async move {
+            let _active_metric = wakey_core::telemetry::ActiveTerminal::new();
             let result = run_terminal(
                 &config,
                 &terminal_id,
@@ -238,6 +244,7 @@ impl TerminalManager {
         Ok(())
     }
 
+    #[tracing::instrument(name = "terminal_close", skip_all, fields(terminal_id = %terminal_id))]
     pub fn close(&self, terminal_id: &TerminalId) -> bool {
         let mut active = self.active.lock().expect("terminal manager poisoned");
         let Some(terminal) = active.get_mut(terminal_id.as_str()) else {
@@ -253,6 +260,7 @@ impl TerminalManager {
         true
     }
 
+    #[tracing::instrument(name = "terminal_resume", skip_all, fields(terminal_id = %terminal_id))]
     pub fn resume(&self, terminal_id: &TerminalId, relay_token: String) -> Result<()> {
         let active = self.active.lock().expect("terminal manager poisoned");
         let session = active
@@ -453,6 +461,7 @@ async fn run_terminal(
                 match read {
                     Ok(0) => break,
                     Ok(count) => {
+                        wakey_core::telemetry::count(wakey_core::telemetry::Counter::TerminalOutputBytes, count as u64);
                         terminal_state.process(&output[..count]);
                         relay_output.queue_live(Message::Binary(output[..count].to_vec().into()));
                     }
@@ -575,6 +584,7 @@ where
                 let Some(message) = incoming else { break; };
                 match message.context("terminal relay websocket receive failed")? {
                     Message::Binary(bytes) => {
+                        wakey_core::telemetry::count(wakey_core::telemetry::Counter::TerminalInputBytes, bytes.len() as u64);
                         pending_input = Some(RelayInput::Binary {
                             generation,
                             bytes: bytes.to_vec(),

wakey-agent/src/tracing.rs

@@ -1,28 +1,19 @@
-use time::UtcOffset;
-use time::format_description::well_known::Rfc3339;
-use tracing_subscriber::fmt::time::OffsetTime;
-use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt};
+use anyhow::Result;
+use wakey_core::telemetry::{TelemetryConfig, TelemetryGuard};
 
-pub fn init(verbose: u8) {
-    let filter = EnvFilter::try_from_default_env()
-        .or_else(|_| EnvFilter::try_new(default_filter(verbose)))
-        .expect("static tracing filter should parse");
-
-    tracing_subscriber::registry()
-        .with(filter)
-        .with(fmt::layer().with_timer(local_offset_timer()))
-        .init();
-}
-
-fn default_filter(verbose: u8) -> &'static str {
-    match verbose {
-        0 => "wakey_agent=info,wakey=info",
-        1 => "wakey_agent=debug,wakey=debug",
-        _ => "wakey_agent=trace,wakey=debug",
-    }
-}
-
-fn local_offset_timer() -> OffsetTime<Rfc3339> {
-    let offset = UtcOffset::current_local_offset().unwrap_or(UtcOffset::UTC);
-    OffsetTime::new(offset, Rfc3339)
+pub fn init(
+    verbose: u8,
+    config: &TelemetryConfig,
+    agent_id: Option<&str>,
+) -> Result<TelemetryGuard> {
+    wakey_core::telemetry::init(
+        config,
+        "wakey-agent",
+        match verbose {
+            0 => "wakey_agent=info,wakey=info,wakey_core=info,opentelemetry=warn",
+            1 => "wakey_agent=debug,wakey=debug,wakey_core=info,opentelemetry=warn",
+            _ => "wakey_agent=trace,wakey=debug,wakey_core=info,opentelemetry=warn",
+        },
+        agent_id,
+    )
 }