fastrace-tower 0.2.0

A tower layer for propagating trace context for fastrace
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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

#![doc = include_str!("../README.md")]

use std::sync::Arc;
use std::task::Context;
use std::task::Poll;

use fastrace::prelude::*;
use http::HeaderValue;
use http::Request;
use tower_layer::Layer;
use tower_service::Service;

/// The standard [W3C Trace Context](https://www.w3.org/TR/trace-context/) header name for passing trace information.
///
/// This is the header key used to propagate trace context between services according to
/// the W3C Trace Context specification.
pub const TRACEPARENT_HEADER: &str = "traceparent";

type SpanContextExtractor =
    Arc<dyn Fn(&http::HeaderMap) -> Option<SpanContext> + Send + Sync + 'static>;

/// Server layer for intercepting and processing trace context in incoming requests.
///
/// This layer extracts tracing context from incoming requests and creates a new span
/// for each request. Add this to your tower server to automatically handle trace context
/// propagation. By default, the layer uses the `traceparent` header to extract a span
/// context and falls back to a random context when the header is missing or invalid.
/// If the configured extractor returns `None`, a noop span is used.
#[derive(Clone)]
pub struct FastraceServerLayer {
    span_context_extractor: SpanContextExtractor,
}

impl Default for FastraceServerLayer {
    fn default() -> Self {
        Self {
            span_context_extractor: Arc::new(|headers| {
                headers
                    .get(TRACEPARENT_HEADER)
                    .and_then(|traceparent| {
                        SpanContext::decode_w3c_traceparent(traceparent.to_str().ok()?)
                    })
                    .or_else(|| Some(SpanContext::random()))
            }),
        }
    }
}

impl FastraceServerLayer {
    /// Configure a custom span context extractor.
    ///
    /// Return `None` to keep the span as noop.
    pub fn with_span_context_extractor<F>(mut self, f: F) -> Self
    where F: Fn(&http::HeaderMap) -> Option<SpanContext> + Send + Sync + 'static {
        self.span_context_extractor = Arc::new(f);
        self
    }
}

impl<S> Layer<S> for FastraceServerLayer {
    type Service = FastraceServerService<S>;

    fn layer(&self, service: S) -> Self::Service {
        FastraceServerService {
            service,
            span_context_extractor: self.span_context_extractor.clone(),
        }
    }
}

/// Server-side service that handles trace context propagation.
///
/// This service extracts trace context from incoming requests and creates
/// spans to track the request processing. It wraps the inner service and augments
/// it with tracing capabilities.
#[derive(Clone)]
pub struct FastraceServerService<S> {
    service: S,
    span_context_extractor: SpanContextExtractor,
}

impl<S, Body> Service<Request<Body>> for FastraceServerService<S>
where S: Service<Request<Body>>
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = fastrace::future::InSpan<S::Future>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(cx)
    }

    fn call(&mut self, req: Request<Body>) -> Self::Future {
        let parent = (self.span_context_extractor)(req.headers());

        let span = if let Some(parent) = parent {
            Span::root(req.method().to_string(), parent)
        } else {
            Span::noop()
        };

        self.service.call(req).in_span(span)
    }
}

/// Client layer for injecting trace context into outgoing requests.
///
/// This layer adds the current trace context to outgoing requests,
/// allowing the receiving service to continue the same trace. Add this
/// to your tower client to automatically propagate trace context.
#[derive(Clone)]
pub struct FastraceClientLayer;

impl<S> Layer<S> for FastraceClientLayer {
    type Service = FastraceClientService<S>;

    fn layer(&self, service: S) -> Self::Service {
        FastraceClientService { service }
    }
}

/// Client-side service that handles trace context propagation.
///
/// This service injects the current trace context into outgoing requests,
/// allowing distributed tracing across service boundaries.
#[derive(Clone)]
pub struct FastraceClientService<S> {
    service: S,
}

impl<S, Body> Service<Request<Body>> for FastraceClientService<S>
where S: Service<Request<Body>>
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(cx)
    }

    fn call(&mut self, mut req: Request<Body>) -> Self::Future {
        if let Some(current) = SpanContext::current_local_parent() {
            req.headers_mut().insert(
                TRACEPARENT_HEADER,
                HeaderValue::from_str(&current.encode_w3c_traceparent()).unwrap(),
            );
        }

        self.service.call(req)
    }
}