mockforge_ws/
ws_tracing.rs

1//! Distributed tracing for WebSocket connections
2
3use mockforge_tracing::{create_request_span, record_error, record_success, Protocol};
4use opentelemetry::{global::BoxedSpan, trace::Span, KeyValue};
5use std::collections::HashMap;
6use tracing::debug;
7
8/// Create a span for a WebSocket connection
9pub fn create_ws_connection_span(path: &str) -> BoxedSpan {
10    create_request_span(
11        Protocol::WebSocket,
12        &format!("WS Connect {}", path),
13        vec![
14            KeyValue::new("ws.path", path.to_string()),
15            KeyValue::new("network.protocol.name", "websocket"),
16        ],
17    )
18}
19
20/// Create a span for a WebSocket message
21pub fn create_ws_message_span(direction: &str, message_type: &str, size: usize) -> BoxedSpan {
22    create_request_span(
23        Protocol::WebSocket,
24        &format!("WS Message {}", direction),
25        vec![
26            KeyValue::new("ws.direction", direction.to_string()),
27            KeyValue::new("ws.message.type", message_type.to_string()),
28            KeyValue::new("ws.message.size", size as i64),
29        ],
30    )
31}
32
33/// Record successful WebSocket connection
34pub fn record_ws_connection_success(
35    span: &mut BoxedSpan,
36    duration_ms: u64,
37    messages_sent: usize,
38    messages_received: usize,
39) {
40    let attributes = vec![
41        KeyValue::new("ws.duration_ms", duration_ms as i64),
42        KeyValue::new("ws.messages.sent", messages_sent as i64),
43        KeyValue::new("ws.messages.received", messages_received as i64),
44    ];
45
46    record_success(span, attributes);
47
48    debug!(
49        duration_ms = duration_ms,
50        messages_sent = messages_sent,
51        messages_received = messages_received,
52        "WebSocket connection completed"
53    );
54}
55
56/// Record WebSocket connection error
57pub fn record_ws_error(span: &mut BoxedSpan, error_message: &str, duration_ms: u64) {
58    span.set_attribute(KeyValue::new("ws.duration_ms", duration_ms as i64));
59    record_error(span, error_message);
60
61    debug!(
62        error_message = error_message,
63        duration_ms = duration_ms,
64        "WebSocket connection failed"
65    );
66}
67
68/// Record WebSocket message success
69pub fn record_ws_message_success(span: &mut BoxedSpan, processing_time_us: u64) {
70    let attributes = vec![KeyValue::new(
71        "ws.processing_time_us",
72        processing_time_us as i64,
73    )];
74
75    record_success(span, attributes);
76}
77
78/// Extract trace context from WebSocket headers
79pub fn extract_ws_trace_context(headers: &HashMap<String, String>) -> opentelemetry::Context {
80    mockforge_tracing::extract_trace_context(headers)
81}
82
83/// Inject trace context into WebSocket response headers
84pub fn inject_ws_trace_context(
85    ctx: &opentelemetry::Context,
86    headers: &mut HashMap<String, String>,
87) {
88    mockforge_tracing::inject_trace_context(ctx, headers);
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use opentelemetry::global;
95    use opentelemetry_sdk::propagation::TraceContextPropagator;
96
97    #[test]
98    fn test_create_ws_connection_span() {
99        global::set_text_map_propagator(TraceContextPropagator::new());
100
101        let mut span = create_ws_connection_span("/ws/chat");
102        record_ws_connection_success(&mut span, 5000, 10, 12);
103    }
104
105    #[test]
106    fn test_create_ws_message_span() {
107        global::set_text_map_propagator(TraceContextPropagator::new());
108
109        let mut span = create_ws_message_span("inbound", "text", 256);
110        record_ws_message_success(&mut span, 150);
111    }
112
113    #[test]
114    fn test_ws_trace_context_propagation() {
115        global::set_text_map_propagator(TraceContextPropagator::new());
116
117        let mut headers = HashMap::new();
118        headers.insert(
119            "traceparent".to_string(),
120            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".to_string(),
121        );
122
123        let ctx = extract_ws_trace_context(&headers);
124
125        let mut output_headers = HashMap::new();
126        inject_ws_trace_context(&ctx, &mut output_headers);
127
128        assert!(output_headers.contains_key("traceparent"));
129    }
130}
mockforge_ws/ws_tracing.rs

mockforge_ws/
ws_tracing.rs
