Trace agent connection phases without long-lived entered spans

This commit is contained in:
lda committed 2026-10-10 02:37:11 +07:00 Verified
1 parent 621cb34a3f
commit 5096205d07
3 files changed
+65 -51

No files matched your search

+6
View File
@@ -94,6 +94,12 @@ existing debug events too; inspect them before enabling it broadly.
and wake dispatch. Counts distinguish grouped, filtered, and returned rows;
wake events identify the chosen agent and whether the route was explicit,
without recording device identifiers or addresses.
- Connectivity uses short `agent_dns`, `agent_connect`, `agent_handshake_send`,
and CC `agent_auth_verify` spans, rather than connection-long spans. Connect
attempts also contribute to operation metrics. Lifecycle events report retries
with backoff and CC disconnect duration. Sending credentials is not an auth ACK;
CC's authentication event is the authoritative result. DNS diagnostics are a
separate lookup, not a measurement of the resolver inside WebSocket connect.
Instrument operation boundaries rather than individual rows or terminal chunks.
Use `skip_all`, explicitly allow count fields, and instrument async futures rather
+49 -42
View File
@@ -4,7 +4,7 @@ use std::net::IpAddr;
use std::time::Instant;
use tokio::time::{Duration, MissedTickBehavior, interval, sleep};
use tokio_tungstenite::{connect_async, tungstenite::Message};
use tracing::{debug, error, info, info_span, warn};
use tracing::{Instrument, debug, error, info, info_span, warn};
use crate::config::AgentConfig;
use crate::dispatch::{dispatch_command, inventory_for_config};
@@ -25,7 +25,7 @@ pub async fn run(config: AgentConfig) -> Result<()> {
backoff = config.reconnect_base_ms.max(100);
}
Err(err) => {
warn!(error = %err, backoff_ms = backoff, "agent session ended; reconnecting");
warn!(agent_id = %config.agent_id, error = %err, backoff_ms = backoff, "agent session ended; reconnecting");
sleep(Duration::from_millis(backoff)).await;
backoff = next_backoff_ms(backoff, config.reconnect_max_ms);
}
@@ -40,20 +40,14 @@ async fn run_once(
crate::terminal::TerminalManagerEvent,
>,
) -> Result<()> {
let session_id = format!(
"{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis())
.unwrap_or_default()
);
let span = info_span!("agent_session", session_id = %session_id, agent_id = %config.agent_id);
let _span_guard = span.enter(); // the ENTIRETY of below is agent_session.
// Connection-long spans export too late and entered guards are unsafe across
// awaits. Short phase spans export while a healthy socket is still running.
let ws_url = websocket_url(&config.server_url)?;
info!(%ws_url, agent_id = %config.agent_id, "connecting agent websocket");
if let Some((dns_resolve_ms, resolved_addrs)) = dns_resolution_diagnostics(&ws_url).await {
if let Some((dns_resolve_ms, resolved_addrs)) = dns_resolution_diagnostics(&ws_url)
.instrument(info_span!("agent_dns", agent_id = %config.agent_id))
.await
{
info!(
host = %ws_url.host_str().unwrap_or(""),
dns_resolve_ms,
@@ -70,9 +64,17 @@ async fn run_once(
}
}
let connect_started = Instant::now();
let (stream, _) = connect_async(ws_url.as_str())
.await
.context("failed to connect websocket")?;
let (stream, _) = async {
let mut operation = wakey_core::telemetry::Operation::new("agent_connect");
let result = connect_async(ws_url.as_str()).await;
if result.is_ok() {
operation.outcome("ok");
}
result
}
.instrument(info_span!("agent_connect", agent_id = %config.agent_id))
.await
.context("failed to connect websocket")?;
let ws_connect_ms = connect_started.elapsed().as_millis() as u64;
info!(%ws_url, agent_id = %config.agent_id, ws_connect_ms, "agent websocket connected");
if ws_connect_ms > 10_000 {
@@ -80,33 +82,38 @@ async fn run_once(
}
let (mut sink, mut source) = stream.split();
send_json(
&mut sink,
&ClientMessage::Hello {
agent_id: config.agent_id.clone(),
capabilities: agent_capabilities(config),
capability_options: AgentCapabilityOptions {
terminal: config
.terminal
.enabled
.then_some(TerminalCapabilityOptions {
max_sessions: config.terminal.max_sessions.max(1),
session_ttl_seconds: config.terminal.session_ttl_seconds,
}),
async {
send_json(
&mut sink,
&ClientMessage::Hello {
agent_id: config.agent_id.clone(),
capabilities: agent_capabilities(config),
capability_options: AgentCapabilityOptions {
terminal: config
.terminal
.enabled
.then_some(TerminalCapabilityOptions {
max_sessions: config.terminal.max_sessions.max(1),
session_ttl_seconds: config.terminal.session_ttl_seconds,
}),
},
},
},
)
)
.await?;
send_json(
&mut sink,
&ClientMessage::Auth {
agent_id: config.agent_id.clone(),
agent_token: config.agent_token.clone(),
},
)
.await?;
send_terminal_sessions(&mut sink, terminal_manager).await?;
Ok::<_, anyhow::Error>(())
}
.instrument(info_span!("agent_handshake_send", agent_id = %config.agent_id))
.await?;
send_json(
&mut sink,
&ClientMessage::Auth {
agent_id: config.agent_id.clone(),
agent_token: config.agent_token.clone(),
},
)
.await?;
send_terminal_sessions(&mut sink, terminal_manager).await?;
info!(agent_id = %config.agent_id, "agent websocket session authenticated");
info!(agent_id = %config.agent_id, "agent hello and credentials sent; awaiting server traffic");
let mut heartbeat = interval(Duration::from_secs(30));
heartbeat.set_missed_tick_behavior(MissedTickBehavior::Skip);
+10 -9
View File
@@ -6,7 +6,7 @@ use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use std::time::Instant;
use tokio::sync::mpsc;
use tracing::{debug, info, info_span, warn};
use tracing::{Instrument, debug, info, info_span, warn};
use uuid::Uuid;
use wakey_agent::protocol::{
AgentCapability, AgentCapabilityOptions, AgentTerminalSession, DEFAULT_TERMINAL_MAX_SESSIONS,
@@ -71,9 +71,9 @@ pub async fn agent_ws(ws: WebSocketUpgrade, State(state): State<AppState>) -> im
async fn handle_agent_socket(state: AppState, socket: WebSocket) {
let connection_id = Uuid::new_v4().to_string();
let span = info_span!("agent_ws_connection", connection_id = %connection_id);
let _span_guard = span.enter();
info!("agent websocket upgraded");
// Lifecycle events are available immediately; do not hold an entered span
// across socket awaits or wait for disconnect to export the connection.
info!(%connection_id, "agent websocket upgraded");
let connected_at = Instant::now();
let (mut write, mut read) = socket.split();
@@ -111,11 +111,11 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
let msg = match frame {
Some(Ok(msg)) => msg,
Some(Err(err)) => {
warn!(error = %err, "agent websocket receive error");
warn!(%connection_id, error = %err, "agent websocket receive error");
break;
}
None => {
info!("agent websocket stream ended by peer");
info!(%connection_id, "agent websocket stream ended by peer");
break;
}
};
@@ -132,14 +132,14 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
)
.await
{
warn!(error = %err, "closing agent websocket due to protocol/auth error");
warn!(%connection_id, error = %err, "closing agent websocket due to protocol/auth error");
break;
}
}
Message::Ping(_) => {}
Message::Pong(_) => {}
Message::Close(frame) => {
info!(close = ?frame, "agent websocket close frame received");
info!(%connection_id, close = ?frame, "agent websocket close frame received");
break;
}
Message::Binary(_) => {
@@ -149,7 +149,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
}
if let Some(agent_id) = connection.authed_agent_id {
info!(agent_id = %agent_id, "agent disconnected");
info!(%connection_id, agent_id = %agent_id, connected_ms = now_duration_ms(connected_at.elapsed()), "agent disconnected");
let mut sessions = state.sessions.write().await;
let should_remove = sessions
.get(&agent_id)
@@ -245,6 +245,7 @@ async fn process_agent_text(
if !state
.store
.verify_agent_token(&agent_id, &agent_token)
.instrument(info_span!("agent_auth_verify", agent_id = %agent_id, %connection_id))
.await
{
warn!(agent_id = %agent_id, "agent auth rejected");