fastrace_tonic/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::task::Context;
4use std::task::Poll;
5
6use fastrace::prelude::*;
7use http::HeaderValue;
8use http::Request;
9use tower_layer::Layer;
10use tower_service::Service;
11
12/// The standard [W3C Trace Context](https://www.w3.org/TR/trace-context/) header name for passing trace information.
13///
14/// This is the header key used to propagate trace context between services according to
15/// the W3C Trace Context specification.
16pub const TRACEPARENT_HEADER: &str = "traceparent";
17
18/// Server layer for intercepting and processing trace context in incoming requests.
19///
20/// This layer extracts tracing context from incoming requests and creates a new span
21/// for each request. Add this to your tonic server to automatically handle trace context
22/// propagation.
23#[derive(Clone)]
24pub struct FastraceServerLayer;
25
26impl<S> Layer<S> for FastraceServerLayer {
27    type Service = FastraceServerService<S>;
28
29    fn layer(&self, service: S) -> Self::Service {
30        FastraceServerService { service }
31    }
32}
33
34/// Server-side service that handles trace context propagation.
35///
36/// This service extracts trace context from incoming requests and creates
37/// spans to track the request processing. It wraps the inner service and augments
38/// it with tracing capabilities.
39#[derive(Clone)]
40pub struct FastraceServerService<S> {
41    service: S,
42}
43
44impl<S, Body> Service<Request<Body>> for FastraceServerService<S>
45where S: Service<Request<Body>>
46{
47    type Response = S::Response;
48    type Error = S::Error;
49    type Future = fastrace::future::InSpan<S::Future>;
50
51    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
52        self.service.poll_ready(cx)
53    }
54
55    fn call(&mut self, req: Request<Body>) -> Self::Future {
56        let headers = req.headers();
57        let parent = headers.get(TRACEPARENT_HEADER).and_then(|traceparent| {
58            SpanContext::decode_w3c_traceparent(traceparent.to_str().ok()?)
59        });
60
61        let span = if let Some(parent) = parent {
62            Span::root(req.uri().to_string(), parent)
63        } else {
64            Span::noop()
65        };
66
67        self.service.call(req).in_span(span)
68    }
69}
70
71/// Client layer for injecting trace context into outgoing requests.
72///
73/// This layer adds the current trace context to outgoing requests,
74/// allowing the receiving service to continue the same trace. Add this
75/// to your tonic client to automatically propagate trace context.
76#[derive(Clone)]
77pub struct FastraceClientLayer;
78
79impl<S> Layer<S> for FastraceClientLayer {
80    type Service = FastraceClientService<S>;
81
82    fn layer(&self, service: S) -> Self::Service {
83        FastraceClientService { service }
84    }
85}
86
87/// Client-side service that handles trace context propagation.
88///
89/// This service injects the current trace context into outgoing requests,
90/// allowing distributed tracing across service boundaries.
91#[derive(Clone)]
92pub struct FastraceClientService<S> {
93    service: S,
94}
95
96impl<S, Body> Service<Request<Body>> for FastraceClientService<S>
97where S: Service<Request<Body>>
98{
99    type Response = S::Response;
100    type Error = S::Error;
101    type Future = S::Future;
102
103    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
104        self.service.poll_ready(cx)
105    }
106
107    fn call(&mut self, mut req: Request<Body>) -> Self::Future {
108        if let Some(current) = SpanContext::current_local_parent() {
109            req.headers_mut().insert(
110                TRACEPARENT_HEADER,
111                HeaderValue::from_str(&current.encode_w3c_traceparent()).unwrap(),
112            );
113        }
114
115        self.service.call(req)
116    }
117}
fastrace_tonic/lib.rs

fastrace_tonic/
lib.rs
