lda/wakey
1
0 Fork 0
Code Issues 6 Pull Requests Actions Packages Projects Releases 20 Wiki Activity

Add shared OTLP telemetry and propagate command traces across agents

Browse Source
This commit is contained in:
lda committed 2026-10-09 23:16:21 +07:00 Verified
1 parent 53c274dbe8
commit 9814754e50
29 files changed
+1227 -213
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

wakey-control-plane/src/api/commands.rs
+13 -8
View File
@@ -4,7 +4,7 @@ use axum::http::StatusCode;
use axum::response::IntoResponse;
use serde::{Deserialize, Serialize};
use std::time::Instant;
use tracing::{info, info_span, warn};
use tracing::{info, warn};
use uuid::Uuid;
use wakey_agent::protocol::{
AgentCapability, AgentCapabilityOptions, AgentCommand, ErrorPayload, RequestId, ServerMessage,
@@ -85,6 +85,7 @@ pub async fn run_command(
}
}
#[tracing::instrument(name = "relay_command", skip_all, fields(agent_id = %agent_id, request_id = tracing::field::Empty, command = tracing::field::Empty))]
pub async fn relay_agent_command(
state: &AppState,
agent_id: &str,
@@ -93,14 +94,11 @@ pub async fn relay_agent_command(
) -> Result<RelayCommandResponse, ApiError> {
let request_id_string = format!("req-{}", Uuid::new_v4());
let command_kind = command_kind(&command);
tracing::Span::current()
.record("request_id", &request_id_string)
.record("command", command_kind);
let started = Instant::now();
let span = info_span!(
"relay_command",
agent_id = %agent_id,
request_id = %request_id_string,
command = %command_kind,
);
let _span_guard = span.enter();
let mut operation = wakey_core::telemetry::Operation::new(command_kind);
let request_id = RequestId::try_from(request_id_string.clone()).map_err(|err| {
ApiError::new(
@@ -149,6 +147,7 @@ pub async fn relay_agent_command(
if let Err(err) = tx.send(SessionEvent::Message(ServerMessage::Command {
request_id,
command,
trace_context: wakey_core::telemetry::current_trace_context(),
})) {
state.pending.lock().await.remove(&request_id_string);
warn!(error = %err, "failed sending command to agent session");
@@ -249,6 +248,7 @@ pub async fn relay_agent_command(
));
}
Err(_) => {
operation.outcome("timeout");
state.pending.lock().await.remove(&request_id_string);
warn!(
timeout_ms = timeout.as_millis() as u64,
@@ -276,6 +276,11 @@ pub async fn relay_agent_command(
}
};
operation.outcome(if response.status == "ok" {
"ok"
} else {
"error"
});
Ok(response)
}
wakey-control-plane/src/api/terminals.rs
+3
View File
@@ -49,6 +49,7 @@ pub struct TerminalSessionResponse {
pub attachment_token: Option<String>,
}
#[tracing::instrument(name = "terminal_create", skip_all)]
pub async fn create_terminal(
State(state): State<AppState>,
Json(request): Json<CreateTerminalRequest>,
@@ -183,6 +184,7 @@ pub async fn list_terminals(State(state): State<AppState>) -> Json<Vec<TerminalS
)
}
#[tracing::instrument(name = "terminal_attach", skip_all)]
pub async fn attach_terminal(
State(state): State<AppState>,
Path(terminal_id): Path<String>,
@@ -210,6 +212,7 @@ pub async fn attach_terminal(
}))
}
#[tracing::instrument(name = "terminal_close", skip_all)]
pub async fn close_terminal(
State(state): State<AppState>,
Path(terminal_id): Path<String>,
wakey-control-plane/src/config/init.rs
+8
View File
@@ -113,6 +113,9 @@ pub fn write_init_config(args: &InitConfigArgs) -> Result<Option<PathBuf>> {
ui_dist_dir,
bootstrap_enroll_tokens,
telemetry: WritableTelemetry {
authorization_env: base_telemetry
.authorization_env
.or_else(|| defaults.telemetry.authorization_env.clone()),
otlp_endpoint: args
.telemetry_otlp_endpoint
.clone()
@@ -223,6 +226,7 @@ bootstrap_enroll_tokens = ["old-token"]
[telemetry]
otlp_endpoint = "http://otel.example"
authorization_env = "WAKEY_OTLP_AUTHORIZATION"
service_name = "old-service"
json_logs = true
"#,
@@ -261,6 +265,10 @@ json_logs = true
Some("/tmp/wakey-base/old.sqlite3")
);
assert_eq!(parsed["telemetry"]["json_logs"].as_bool(), Some(true));
assert_eq!(
parsed["telemetry"]["authorization_env"].as_str(),
Some("WAKEY_OTLP_AUTHORIZATION")
);
assert_eq!(
parsed["telemetry"]["service_name"].as_str(),
Some("old-service")
wakey-control-plane/src/config/resolve.rs
+1
View File
@@ -231,6 +231,7 @@ fn resolve_admin_public_url(
fn resolve_telemetry(file: Option<FileTelemetryConfig>) -> TelemetryConfig {
let mut out = TelemetryConfig::default();
if let Some(file) = file {
out.authorization_env = file.authorization_env;
if let Some(endpoint) = file.otlp_endpoint {
let trimmed = endpoint.trim().to_string();
if !trimmed.is_empty() {
wakey-control-plane/src/config/types.rs
+6
View File
@@ -23,6 +23,7 @@ pub struct DaemonConfig {
#[derive(Debug, Clone)]
pub struct TelemetryConfig {
pub otlp_endpoint: Option<String>,
pub authorization_env: Option<String>,
pub service_name: String,
pub json_logs: bool,
}
@@ -32,6 +33,7 @@ impl Default for TelemetryConfig {
let defaults = &DEFAULT_CONFIG.telemetry;
Self {
otlp_endpoint: defaults.otlp_endpoint.clone(),
authorization_env: defaults.authorization_env.clone(),
service_name: defaults.service_name.clone(),
json_logs: defaults.json_logs,
}
@@ -57,6 +59,7 @@ pub(crate) struct FileConfig {
#[derive(Debug, Deserialize, Default)]
pub(crate) struct FileTelemetryConfig {
pub(crate) otlp_endpoint: Option<String>,
pub(crate) authorization_env: Option<String>,
pub(crate) service_name: Option<String>,
pub(crate) json_logs: Option<bool>,
}
@@ -81,6 +84,8 @@ pub(crate) struct WritableConfig {
pub(crate) struct WritableTelemetry {
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) otlp_endpoint: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) authorization_env: Option<String>,
pub(crate) service_name: String,
pub(crate) json_logs: bool,
}
@@ -98,6 +103,7 @@ pub(crate) static DEFAULT_CONFIG: LazyLock<WritableConfig> = LazyLock::new(|| Wr
bootstrap_enroll_tokens: Vec::new(),
telemetry: WritableTelemetry {
otlp_endpoint: None,
authorization_env: None,
service_name: "wakey-control-plane".to_string(),
json_logs: false,
},
wakey-control-plane/src/main.rs
+4 -2
View File
@@ -28,8 +28,10 @@ async fn main() -> Result<()> {
}
}
let daemon = config::DaemonConfig::from_serve_args(&args)?;
tracing::init(cli.verbose, &daemon.telemetry)?;
runtime::serve(daemon).await
let telemetry = tracing::init(cli.verbose, &daemon.telemetry)?;
let result = runtime::serve(daemon).await;
tokio::task::spawn_blocking(move || telemetry.shutdown()).await?;
result
}
Command::InitConfig(args) => {
if let Some(path) = config::write_init_config(&args)? {
wakey-control-plane/src/runtime/mod.rs
+7
View File
@@ -213,6 +213,8 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
use tokio::signal::unix::{SignalKind, signal};
let mut server = server;
let mut hup = signal(SignalKind::hangup()).context("failed to install SIGHUP handler")?;
let mut terminate =
signal(SignalKind::terminate()).context("failed to install SIGTERM handler")?;
let mut gc_tick = tokio::time::interval(Duration::from_secs(300));
gc_tick.set_missed_tick_behavior(MissedTickBehavior::Skip);
@@ -224,6 +226,11 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
server.abort();
break;
}
_ = terminate.recv() => {
info!("SIGTERM received; shutting down control-plane");
server.abort();
break;
}
_ = hup.recv() => {
match app_state.store.reload_from_disk().await {
Ok(()) => info!("reloaded state from disk"),
wakey-control-plane/src/tracing.rs
+18 -95
View File
@@ -1,99 +1,22 @@
use anyhow::{Context, Result};
use opentelemetry::global;
use opentelemetry::trace::TracerProvider as _;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::Resource;
use opentelemetry_sdk::trace::{SdkTracerProvider, Tracer};
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 crate::config::TelemetryConfig;
use anyhow::Result;
pub use wakey_core::telemetry::TelemetryGuard;
pub fn init(verbose: u8, telemetry: &TelemetryConfig) -> Result<()> {
let filter = EnvFilter::try_from_default_env()
.or_else(|_| EnvFilter::try_new(default_filter(verbose)))
.expect("static tracing filter should parse");
let otel = build_otel_layer(telemetry)?;
if telemetry.json_logs {
if let Some(otel_layer) = otel {
tracing_subscriber::registry()
.with(otel_layer)
.with(filter)
.with(fmt::layer().json().with_timer(local_offset_timer()))
.init();
tracing::info!(endpoint = ?telemetry.otlp_endpoint, json_logs = telemetry.json_logs, "tracing initialized with otlp exporter");
} else {
tracing_subscriber::registry()
.with(filter)
.with(fmt::layer().json().with_timer(local_offset_timer()))
.init();
tracing::info!(
json_logs = telemetry.json_logs,
"tracing initialized without otlp exporter"
);
}
} else if let Some(otel_layer) = otel {
tracing_subscriber::registry()
.with(otel_layer)
.with(filter)
.with(fmt::layer().with_timer(local_offset_timer()))
.init();
tracing::info!(endpoint = ?telemetry.otlp_endpoint, json_logs = telemetry.json_logs, "tracing initialized with otlp exporter");
} else {
tracing_subscriber::registry()
.with(filter)
.with(fmt::layer().with_timer(local_offset_timer()))
.init();
tracing::info!(
json_logs = telemetry.json_logs,
"tracing initialized without otlp exporter"
);
}
Ok(())
}
fn build_otel_layer(
telemetry: &TelemetryConfig,
) -> Result<Option<tracing_opentelemetry::OpenTelemetryLayer<tracing_subscriber::Registry, Tracer>>>
{
let Some(endpoint) = telemetry.otlp_endpoint.as_deref() else {
return Ok(None);
pub fn init(verbose: u8, telemetry: &TelemetryConfig) -> Result<TelemetryGuard> {
let config = wakey_core::telemetry::TelemetryConfig {
otlp_endpoint: telemetry.otlp_endpoint.clone(),
authorization_env: telemetry.authorization_env.clone(),
service_name: telemetry.service_name.clone(),
json_logs: telemetry.json_logs,
};
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_tonic()
.with_endpoint(endpoint)
.build()
.context("failed to build OTLP span exporter")?;
let provider = SdkTracerProvider::builder()
.with_batch_exporter(exporter)
.with_resource(
Resource::builder_empty()
.with_service_name(telemetry.service_name.clone())
.build(),
)
.build();
let tracer = provider.tracer(telemetry.service_name.clone());
global::set_tracer_provider(provider);
Ok(Some(tracing_opentelemetry::layer().with_tracer(tracer)))
}
fn default_filter(verbose: u8) -> &'static str {
match verbose {
0 => "wakey_control_plane=info",
1 => "wakey_control_plane=debug",
_ => "wakey_control_plane=trace",
}
}
fn local_offset_timer() -> OffsetTime<Rfc3339> {
let offset = UtcOffset::current_local_offset().unwrap_or(UtcOffset::UTC);
OffsetTime::new(offset, Rfc3339)
wakey_core::telemetry::init(
&config,
"wakey-control-plane",
match verbose {
0 => "wakey_control_plane=info,wakey_core=info,opentelemetry=warn",
1 => "wakey_control_plane=debug,wakey_core=info,opentelemetry=warn",
_ => "wakey_control_plane=trace,wakey_core=info,opentelemetry=warn",
},
None,
)
}
wakey-control-plane/src/ws.rs
+10 -4
View File
@@ -309,12 +309,18 @@ async fn process_agent_text(
anyhow::bail!("device_snapshot for unauthenticated or mismatched agent");
}
ensure_current_session(state, &agent_id, connection_id).await?;
match state
.store
.replace_agent_device_snapshot(&agent_id, &devices)
.await
let mut operation = wakey_core::telemetry::Operation::new("snapshot_store");
let span = tracing::info_span!("store_device_snapshot", agent_id = %agent_id, devices = devices.len());
match tracing::Instrument::instrument(
state
.store
.replace_agent_device_snapshot(&agent_id, &devices),
span,
)
.await
{
Ok(accepted) => {
operation.outcome("ok");
debug!(agent_id = %agent_id, accepted, "agent websocket device snapshot accepted");
}
Err(err) => {