mockforge-grpc 0.3.127

gRPC protocol support for MockForge
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

//! Distributed tracing for gRPC reflection proxy

use mockforge_tracing::{create_request_span, record_error, record_success, Protocol};
use opentelemetry::{global::BoxedSpan, trace::Span, KeyValue};
use std::collections::HashMap;
use tracing::debug;

/// Create a span for a gRPC method invocation
pub fn create_grpc_span(service_name: &str, method_name: &str) -> BoxedSpan {
    let method_full = format!("{}::{}", service_name, method_name);

    create_request_span(
        Protocol::Grpc,
        &method_full,
        vec![
            KeyValue::new("rpc.system", "grpc"),
            KeyValue::new("rpc.service", service_name.to_string()),
            KeyValue::new("rpc.method", method_name.to_string()),
        ],
    )
}

/// Record successful gRPC method completion
pub fn record_grpc_success(
    span: &mut BoxedSpan,
    duration_ms: u64,
    request_size: Option<usize>,
    response_size: Option<usize>,
) {
    let mut attributes = vec![
        KeyValue::new("rpc.grpc.status_code", 0i64), // OK
        KeyValue::new("rpc.duration_ms", duration_ms as i64),
    ];

    if let Some(size) = request_size {
        attributes.push(KeyValue::new("rpc.request.size", size as i64));
    }

    if let Some(size) = response_size {
        attributes.push(KeyValue::new("rpc.response.size", size as i64));
    }

    record_success(span, attributes);

    debug!(duration_ms = duration_ms, "gRPC method completed successfully");
}

/// Record gRPC method error
pub fn record_grpc_error(
    span: &mut BoxedSpan,
    error_code: i32,
    error_message: &str,
    duration_ms: u64,
) {
    span.set_attribute(KeyValue::new("rpc.grpc.status_code", error_code as i64));
    span.set_attribute(KeyValue::new("rpc.duration_ms", duration_ms as i64));

    record_error(span, error_message);

    debug!(
        error_code = error_code,
        error_message = error_message,
        duration_ms = duration_ms,
        "gRPC method failed"
    );
}

/// Extract trace context from gRPC metadata
pub fn extract_grpc_trace_context(metadata: &HashMap<String, String>) -> opentelemetry::Context {
    mockforge_tracing::extract_trace_context(metadata)
}

/// Inject trace context into gRPC metadata
pub fn inject_grpc_trace_context(
    ctx: &opentelemetry::Context,
    metadata: &mut HashMap<String, String>,
) {
    mockforge_tracing::inject_trace_context(ctx, metadata);
}

#[cfg(test)]
mod tests {
    use super::*;
    use opentelemetry::global;
    use opentelemetry_sdk::propagation::TraceContextPropagator;

    #[test]
    fn test_create_grpc_span() {
        global::set_text_map_propagator(TraceContextPropagator::new());

        let mut span = create_grpc_span("UserService", "GetUser");
        record_grpc_success(&mut span, 42, Some(100), Some(200));

        // Span should be created successfully
    }

    #[test]
    fn test_trace_context_propagation() {
        global::set_text_map_propagator(TraceContextPropagator::new());

        let mut metadata = HashMap::new();
        metadata.insert(
            "traceparent".to_string(),
            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".to_string(),
        );

        let ctx = extract_grpc_trace_context(&metadata);

        let mut output_metadata = HashMap::new();
        inject_grpc_trace_context(&ctx, &mut output_metadata);

        // Context should be preserved
        assert!(output_metadata.contains_key("traceparent"));
    }
}