featherbit/plugins/native/
zipkin.rs

1//! Distributed tracing via Zipkin / B3 propagation (`zipkin`).
2//!
3//! The featherbit port of Apache APISIX's `zipkin` plugin. Where APISIX runs it
4//! across nginx phases, featherbit models it as a **start/end node pair** wired
5//! around the `upstream` node, using the shared
6//! [`trace`](crate::plugins::util::trace) helper to carry a [`SpanContext`]
7//! through `context.message`:
8//!
9//! - **start node** (`phase: start`, placed right after `listener`, before
10//!   `upstream`): extracts B3 propagation from either the single `b3` header or
11//!   the `x-b3-*` multi-header form. If present, this hop *continues* the
12//!   incoming trace (same `trace_id`, the caller's span becomes our parent, the
13//!   incoming sampled flag is honored). If absent, a *new* trace is started
14//!   (`new_trace_id()`, no parent, sampled per `sample_ratio`). This hop gets a
15//!   fresh `span_id` and start time, the span is stored via [`store_span`], and
16//!   the outgoing `x-b3-*` headers are injected so the upstream continues the
17//!   trace.
18//! - **end node** (`phase: end`, placed after `upstream`, right before
19//!   `client`): loads the span, computes its duration, and — when sampled —
20//!   builds a Zipkin v2 JSON span array and **fire-and-forgets** a POST to the
21//!   collector `endpoint` on a detached `tokio::spawn` task (the same
22//!   best-effort pattern as `proxy-mirror`). The export result is ignored and
23//!   never blocks or fails the request.
24//!
25//! Both nodes always return `Ok`: tracing is observability, so it never
26//! short-circuits the request through an error port. If the end node finds no
27//! stored span it passes through untouched.
28//!
29//! ## Deviations from APISIX
30//! - Spans are exported one-per-request (fire-and-forget) rather than through a
31//!   batched reporter; only span_version 2 (the APISIX default) is modeled, and
32//!   only a single `SERVER` span per hop (no `apisix.rewrite`/`proxy` children).
33//! - The new-trace sampling draw is a *pseudo-random*, per-trace-consistent
34//!   hash of the trace id (no `rand` crate), so a given trace id always samples
35//!   the same way in a process.
36
37use std::collections::HashMap;
38use std::sync::Arc;
39use std::time::Duration;
40
41use async_trait::async_trait;
42use serde_json::{json, Value};
43
44use crate::context::Context;
45use crate::outbound::{OutboundClient, OutboundRequest};
46use crate::plugins::resources::PluginResources;
47use crate::plugins::util::trace::{
48    build_b3_headers, load_span, new_span_id, new_trace_id, now_ms, parse_b3, store_span,
49    SpanContext,
50};
51use crate::plugins::{Plugin, PluginOutput, PluginResult};
52
53/// Which node of the tracing pair this instance is.
54#[derive(Debug, Clone, Copy, PartialEq)]
55enum Phase {
56    /// Before `upstream`: extract/create the span and inject `x-b3-*` headers.
57    Start,
58    /// After `upstream`: export the finished span (fire-and-forget).
59    End,
60}
61
62/// A pseudo-random draw in `[0.0, 1.0)` derived deterministically from a trace
63/// id, so the same trace id always samples the same way within a process.
64/// Not cryptographic; good enough for `sample_ratio` without a `rand` crate.
65fn ratio_draw(trace_id: &str) -> f64 {
66    use std::collections::hash_map::DefaultHasher;
67    use std::hash::{Hash, Hasher};
68    let mut h = DefaultHasher::new();
69    trace_id.hash(&mut h);
70    h.finish() as f64 / (u64::MAX as f64 + 1.0)
71}
72
73/// One node of the Zipkin tracing pair.
74pub struct ZipkinPlugin {
75    phase: Phase,
76    /// Full collector URL spans POST to, e.g.
77    /// `http://localhost:9411/api/v2/spans`.
78    endpoint: String,
79    service_name: String,
80    /// Optional `localEndpoint.ipv4` reported on each span.
81    server_addr: Option<String>,
82    /// Fraction of new traces to sample, in `0.0..=1.0`.
83    sample_ratio: f64,
84    ssl_verify: bool,
85    timeout: Duration,
86    client: Arc<OutboundClient>,
87}
88
89impl ZipkinPlugin {
90    /// Builds the plugin from node config.
91    ///
92    /// Config keys:
93    ///
94    /// | Key | Type | Default | Description |
95    /// |---|---|---|---|
96    /// | `phase` | string | — (**required**) | `start` (before upstream) or `end` (after upstream). |
97    /// | `endpoint` | string | — (**required**) | Full Zipkin v2 collector URL, e.g. `http://localhost:9411/api/v2/spans`. |
98    /// | `service_name` | string | `"featherbit"` | `localEndpoint.serviceName` reported on each span. |
99    /// | `server_addr` | string | — | Optional `localEndpoint.ipv4` for the reporting host. |
100    /// | `sample_ratio` | number | `1.0` | Fraction of new traces to sample, in `0.0..=1.0`. |
101    /// | `ssl_verify` | bool | `true` | Verify the collector's TLS certificate. |
102    /// | `timeout` | int (seconds) | `3` | Export request timeout (end node). |
103    ///
104    /// Fails fast on an unknown `phase`, a missing/empty `endpoint`, or a
105    /// `sample_ratio` outside `0.0..=1.0`. The same config (including
106    /// `endpoint`) is placed on both the start and end nodes.
107    ///
108    /// ```yaml
109    /// # start node — right after listener, before upstream
110    /// - id: zipkin-start
111    ///   type: zipkin
112    ///   config:
113    ///     phase: start
114    ///     endpoint: http://localhost:9411/api/v2/spans
115    ///     service_name: my-gateway
116    ///     sample_ratio: 0.1
117    /// # end node — after upstream, right before client
118    /// - id: zipkin-end
119    ///   type: zipkin
120    ///   config:
121    ///     phase: end
122    ///     endpoint: http://localhost:9411/api/v2/spans
123    /// ```
124    pub fn from_config(
125        config: &HashMap<String, Value>,
126        resources: &Arc<PluginResources>,
127    ) -> Result<Self, String> {
128        let phase = match config.get("phase").and_then(|v| v.as_str()) {
129            Some("start") => Phase::Start,
130            Some("end") => Phase::End,
131            Some(other) => {
132                return Err(format!(
133                    "zipkin `phase` must be `start` or `end` (got `{other}`)"
134                ))
135            }
136            None => return Err("zipkin requires `phase` (`start` or `end`)".to_string()),
137        };
138
139        let endpoint = config
140            .get("endpoint")
141            .and_then(|v| v.as_str())
142            .filter(|s| !s.is_empty())
143            .ok_or("zipkin requires `endpoint` (e.g. \"http://localhost:9411/api/v2/spans\")")?
144            .to_string();
145
146        let service_name = config
147            .get("service_name")
148            .and_then(|v| v.as_str())
149            .unwrap_or("featherbit")
150            .to_string();
151
152        let server_addr = config
153            .get("server_addr")
154            .and_then(|v| v.as_str())
155            .filter(|s| !s.is_empty())
156            .map(String::from);
157
158        let sample_ratio = match config.get("sample_ratio") {
159            None => 1.0,
160            Some(v) => v
161                .as_f64()
162                .filter(|r| (0.0..=1.0).contains(r))
163                .ok_or("zipkin `sample_ratio` must be a number in 0.0..=1.0")?,
164        };
165
166        let ssl_verify = config
167            .get("ssl_verify")
168            .and_then(|v| v.as_bool())
169            .unwrap_or(true);
170        let timeout =
171            Duration::from_secs(config.get("timeout").and_then(|v| v.as_u64()).unwrap_or(3));
172
173        Ok(Self {
174            phase,
175            endpoint,
176            service_name,
177            server_addr,
178            sample_ratio,
179            ssl_verify,
180            timeout,
181            client: resources.outbound.clone(),
182        })
183    }
184
185    /// Decides whether a brand-new trace should be sampled, per `sample_ratio`.
186    fn sample_new(&self, trace_id: &str) -> bool {
187        if self.sample_ratio >= 1.0 {
188            true
189        } else if self.sample_ratio <= 0.0 {
190            false
191        } else {
192            ratio_draw(trace_id) < self.sample_ratio
193        }
194    }
195
196    /// START: extract or create the span, store it, and inject `x-b3-*`.
197    fn run_start(&self, ctx: &mut Context) {
198        let first = |name: &str| {
199            ctx.request
200                .headers
201                .get(name)
202                .and_then(|v| v.first())
203                .map(|s| s.as_str())
204        };
205        let incoming = parse_b3(
206            first("b3"),
207            first("x-b3-traceid"),
208            first("x-b3-spanid"),
209            first("x-b3-sampled"),
210        );
211
212        let (trace_id, parent_span_id, sampled) = match incoming {
213            Some((trace_id, parent, sampled)) => (trace_id, Some(parent), sampled),
214            None => {
215                let trace_id = new_trace_id();
216                let sampled = self.sample_new(&trace_id);
217                (trace_id, None, sampled)
218            }
219        };
220
221        let span = SpanContext {
222            trace_id,
223            span_id: new_span_id(),
224            parent_span_id,
225            sampled,
226            start_ms: now_ms(),
227        };
228
229        // Inject the downstream B3 headers so the upstream continues the trace.
230        for (name, value) in build_b3_headers(&span) {
231            ctx.request.headers.insert(name, vec![value]);
232        }
233
234        store_span(ctx, &span);
235    }
236
237    /// END: load the span and fire-and-forget the Zipkin export when sampled.
238    fn run_end(&self, ctx: &Context) {
239        let Some(span) = load_span(ctx) else {
240            return;
241        };
242        if !span.sampled {
243            return;
244        }
245
246        let payload = build_zipkin(&span, ctx, &self.service_name, self.server_addr.as_deref());
247        let body = match serde_json::to_vec(&payload) {
248            Ok(b) => b,
249            Err(_) => return,
250        };
251
252        let req = OutboundRequest {
253            method: http::Method::POST,
254            url: self.endpoint.clone(),
255            headers: vec![("Content-Type".to_string(), "application/json".to_string())],
256            body: body.into(),
257            timeout: self.timeout,
258            ssl_verify: self.ssl_verify,
259            tls: None,
260        };
261        let client = self.client.clone();
262        // Fire-and-forget: the export never blocks or affects the request path.
263        tokio::spawn(async move {
264            let _ = client.request(req).await;
265        });
266    }
267}
268
269/// Builds the Zipkin v2 JSON span array for a finished span (pure; no I/O).
270///
271/// `timestamp` and `duration` are microseconds (Zipkin's unit); `parentId` is
272/// omitted for a root span.
273fn build_zipkin(
274    span: &SpanContext,
275    ctx: &Context,
276    service_name: &str,
277    server_addr: Option<&str>,
278) -> Value {
279    let timestamp_micros = span.start_ms.saturating_mul(1_000);
280    let duration_micros = span.duration_ms().saturating_mul(1_000);
281    let name = format!("{} {}", ctx.request.method, ctx.request.path);
282
283    let mut local_endpoint = json!({ "serviceName": service_name });
284    if let Some(addr) = server_addr {
285        local_endpoint["ipv4"] = json!(addr);
286    }
287
288    let mut span_json = json!({
289        "traceId": span.trace_id,
290        "id": span.span_id,
291        "name": name,
292        "kind": "SERVER",
293        "timestamp": timestamp_micros,
294        "duration": duration_micros,
295        "localEndpoint": local_endpoint,
296        "tags": {
297            "http.method": ctx.request.method,
298            "http.status_code": ctx.response.status_code.to_string(),
299            "http.path": ctx.request.path,
300        },
301    });
302    if let Some(parent) = &span.parent_span_id {
303        span_json["parentId"] = json!(parent);
304    }
305
306    json!([span_json])
307}
308
309#[async_trait]
310impl Plugin for ZipkinPlugin {
311    fn plugin_type(&self) -> &str {
312        "zipkin"
313    }
314
315    fn reads_response_body(&self) -> bool {
316        false
317    }
318
319    async fn execute(&self, mut ctx: Context) -> PluginResult {
320        match self.phase {
321            Phase::Start => self.run_start(&mut ctx),
322            Phase::End => self.run_end(&ctx),
323        }
324        Ok(PluginOutput::success(ctx))
325    }
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331    use crate::context::{GatewayRequest, GatewayResponse, Protocol};
332    use bytes::Bytes;
333
334    fn plugin(config: Value) -> Result<ZipkinPlugin, String> {
335        let map: HashMap<String, Value> = serde_json::from_value(config).unwrap();
336        ZipkinPlugin::from_config(&map, &PluginResources::empty())
337    }
338
339    fn test_ctx() -> Context {
340        Context {
341            request: GatewayRequest {
342                method: "POST".to_string(),
343                path: "/api/users".to_string(),
344                host: "example.com".to_string(),
345                scheme: "http".to_string(),
346                headers: HashMap::new(),
347                query_params: HashMap::new(),
348                body: Bytes::new(),
349                remote_addr: "10.0.0.1:5555".to_string(),
350                protocol: Protocol::Http1,
351            },
352            response: GatewayResponse {
353                status_code: 201,
354                headers: HashMap::new(),
355                body: Bytes::new(),
356                stream: None,
357            },
358            message: HashMap::new(),
359            errors: Vec::new(),
360        }
361    }
362
363    fn a_span(sampled: bool, parent: Option<&str>) -> SpanContext {
364        SpanContext {
365            trace_id: "0af7651916cd43dd8448eb211c80319c".to_string(),
366            span_id: "b7ad6b7169203331".to_string(),
367            parent_span_id: parent.map(String::from),
368            sampled,
369            start_ms: 1_700_000_000_000,
370        }
371    }
372
373    #[test]
374    fn config_requires_phase_and_endpoint() {
375        assert!(plugin(json!({})).is_err());
376        assert!(plugin(json!({ "phase": "start" })).is_err()); // missing endpoint
377        assert!(plugin(json!({ "phase": "bogus", "endpoint": "http://z" })).is_err());
378        assert!(plugin(json!({ "phase": "start", "endpoint": "http://z" })).is_ok());
379        // sample_ratio out of range
380        assert!(
381            plugin(json!({ "phase": "start", "endpoint": "http://z", "sample_ratio": 2 })).is_err()
382        );
383    }
384
385    #[test]
386    fn sample_ratio_extremes() {
387        let never =
388            plugin(json!({ "phase": "start", "endpoint": "http://z", "sample_ratio": 0 })).unwrap();
389        let always =
390            plugin(json!({ "phase": "start", "endpoint": "http://z", "sample_ratio": 1 })).unwrap();
391        for _ in 0..50 {
392            assert!(!never.sample_new("abc"));
393            assert!(always.sample_new("abc"));
394        }
395    }
396
397    #[test]
398    fn zipkin_payload_shape() {
399        let span = a_span(true, Some("0020000000000001"));
400        let v = build_zipkin(&span, &test_ctx(), "svc", Some("10.0.0.5"));
401        let s = &v[0];
402        assert_eq!(s["traceId"], "0af7651916cd43dd8448eb211c80319c");
403        assert_eq!(s["id"], "b7ad6b7169203331");
404        assert_eq!(s["parentId"], "0020000000000001");
405        assert_eq!(s["kind"], "SERVER");
406        assert_eq!(s["name"], "POST /api/users");
407        // micros
408        assert_eq!(s["timestamp"], 1_700_000_000_000_000u64);
409        assert_eq!(s["localEndpoint"]["serviceName"], "svc");
410        assert_eq!(s["localEndpoint"]["ipv4"], "10.0.0.5");
411        assert_eq!(s["tags"]["http.method"], "POST");
412        assert_eq!(s["tags"]["http.status_code"], "201");
413        assert_eq!(s["tags"]["http.path"], "/api/users");
414    }
415
416    #[test]
417    fn zipkin_omits_parent_when_root() {
418        let span = a_span(true, None);
419        let v = build_zipkin(&span, &test_ctx(), "svc", None);
420        assert!(v[0].get("parentId").is_none());
421        assert!(v[0]["localEndpoint"].get("ipv4").is_none());
422    }
423
424    #[tokio::test]
425    async fn start_creates_new_trace_and_injects_headers() {
426        let p =
427            plugin(json!({ "phase": "start", "endpoint": "http://z", "sample_ratio": 1 })).unwrap();
428        let out = p.execute(test_ctx()).await.unwrap();
429        let span = load_span(&out.context).expect("span stored");
430        assert!(span.parent_span_id.is_none());
431        assert!(span.sampled);
432        assert_eq!(span.trace_id.len(), 32);
433        // downstream B3 headers injected, matching the span
434        assert_eq!(
435            out.context.request.headers["x-b3-traceid"][0],
436            span.trace_id
437        );
438        assert_eq!(out.context.request.headers["x-b3-spanid"][0], span.span_id);
439        assert_eq!(out.context.request.headers["x-b3-sampled"][0], "1");
440    }
441
442    #[tokio::test]
443    async fn start_continues_incoming_b3_single() {
444        let p = plugin(json!({ "phase": "start", "endpoint": "http://z" })).unwrap();
445        let mut ctx = test_ctx();
446        ctx.request
447            .headers
448            .insert("b3".to_string(), vec!["trace1-span1-1".to_string()]);
449        let out = p.execute(ctx).await.unwrap();
450        let span = load_span(&out.context).unwrap();
451        assert_eq!(span.trace_id, "trace1");
452        assert_eq!(span.parent_span_id.as_deref(), Some("span1"));
453        assert!(span.sampled);
454        // our own span id replaces the caller's
455        assert_ne!(span.span_id, "span1");
456    }
457
458    #[tokio::test]
459    async fn start_continues_incoming_b3_multi() {
460        let p = plugin(json!({ "phase": "start", "endpoint": "http://z" })).unwrap();
461        let mut ctx = test_ctx();
462        ctx.request
463            .headers
464            .insert("x-b3-traceid".to_string(), vec!["t9".to_string()]);
465        ctx.request
466            .headers
467            .insert("x-b3-spanid".to_string(), vec!["s9".to_string()]);
468        ctx.request
469            .headers
470            .insert("x-b3-sampled".to_string(), vec!["0".to_string()]);
471        let out = p.execute(ctx).await.unwrap();
472        let span = load_span(&out.context).unwrap();
473        assert_eq!(span.trace_id, "t9");
474        assert_eq!(span.parent_span_id.as_deref(), Some("s9"));
475        assert!(!span.sampled);
476    }
477
478    #[tokio::test]
479    async fn end_passes_through_without_span() {
480        let p = plugin(json!({ "phase": "end", "endpoint": "http://z" })).unwrap();
481        let out = p.execute(test_ctx()).await.unwrap();
482        assert_eq!(out.context.response.status_code, 201);
483    }
484
485    #[tokio::test]
486    async fn end_returns_ok_with_stored_span() {
487        let p = plugin(json!({ "phase": "end", "endpoint": "http://127.0.0.1:1/api/v2/spans" }))
488            .unwrap();
489        let mut ctx = test_ctx();
490        store_span(&mut ctx, &a_span(true, None));
491        // The spawned export is best-effort; we only assert execute returns Ok.
492        let out = p.execute(ctx).await.unwrap();
493        assert_eq!(out.context.request.path, "/api/users");
494    }
495
496    #[tokio::test]
497    async fn end_skips_export_when_unsampled() {
498        let p = plugin(json!({ "phase": "end", "endpoint": "http://z" })).unwrap();
499        let mut ctx = test_ctx();
500        store_span(&mut ctx, &a_span(false, None));
501        let out = p.execute(ctx).await.unwrap();
502        assert_eq!(out.context.response.status_code, 201);
503    }
504}
featherbit/plugins/native/zipkin.rs

featherbit/plugins/native/
zipkin.rs
