1use 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#[derive(Debug, Clone, Copy, PartialEq)]
55enum Phase {
56 Start,
58 End,
60}
61
62fn 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
73pub struct ZipkinPlugin {
75 phase: Phase,
76 endpoint: String,
79 service_name: String,
80 server_addr: Option<String>,
82 sample_ratio: f64,
84 ssl_verify: bool,
85 timeout: Duration,
86 client: Arc<OutboundClient>,
87}
88
89impl ZipkinPlugin {
90 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 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 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 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 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 tokio::spawn(async move {
264 let _ = client.request(req).await;
265 });
266 }
267}
268
269fn 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 async fn execute(
316 &self,
317 mut ctx: Context,
318 _named_inputs: &HashMap<String, Value>,
319 ) -> PluginResult {
320 match self.phase {
321 Phase::Start => self.run_start(&mut ctx),
322 Phase::End => self.run_end(&ctx),
323 }
324 Ok(PluginOutput {
325 context: ctx,
326 named_outputs: HashMap::new(),
327 })
328 }
329}
330
331#[cfg(test)]
332mod tests {
333 use super::*;
334 use crate::context::{GatewayRequest, GatewayResponse, Protocol};
335 use bytes::Bytes;
336
337 fn plugin(config: Value) -> Result<ZipkinPlugin, String> {
338 let map: HashMap<String, Value> = serde_json::from_value(config).unwrap();
339 ZipkinPlugin::from_config(&map, &PluginResources::empty())
340 }
341
342 fn test_ctx() -> Context {
343 Context {
344 request: GatewayRequest {
345 method: "POST".to_string(),
346 path: "/api/users".to_string(),
347 host: "example.com".to_string(),
348 scheme: "http".to_string(),
349 headers: HashMap::new(),
350 query_params: HashMap::new(),
351 body: Bytes::new(),
352 remote_addr: "10.0.0.1:5555".to_string(),
353 protocol: Protocol::Http1,
354 },
355 response: GatewayResponse {
356 status_code: 201,
357 headers: HashMap::new(),
358 body: Bytes::new(),
359 },
360 message: HashMap::new(),
361 errors: Vec::new(),
362 }
363 }
364
365 fn a_span(sampled: bool, parent: Option<&str>) -> SpanContext {
366 SpanContext {
367 trace_id: "0af7651916cd43dd8448eb211c80319c".to_string(),
368 span_id: "b7ad6b7169203331".to_string(),
369 parent_span_id: parent.map(String::from),
370 sampled,
371 start_ms: 1_700_000_000_000,
372 }
373 }
374
375 #[test]
376 fn config_requires_phase_and_endpoint() {
377 assert!(plugin(json!({})).is_err());
378 assert!(plugin(json!({ "phase": "start" })).is_err()); assert!(plugin(json!({ "phase": "bogus", "endpoint": "http://z" })).is_err());
380 assert!(plugin(json!({ "phase": "start", "endpoint": "http://z" })).is_ok());
381 assert!(
383 plugin(json!({ "phase": "start", "endpoint": "http://z", "sample_ratio": 2 })).is_err()
384 );
385 }
386
387 #[test]
388 fn sample_ratio_extremes() {
389 let never =
390 plugin(json!({ "phase": "start", "endpoint": "http://z", "sample_ratio": 0 })).unwrap();
391 let always =
392 plugin(json!({ "phase": "start", "endpoint": "http://z", "sample_ratio": 1 })).unwrap();
393 for _ in 0..50 {
394 assert!(!never.sample_new("abc"));
395 assert!(always.sample_new("abc"));
396 }
397 }
398
399 #[test]
400 fn zipkin_payload_shape() {
401 let span = a_span(true, Some("0020000000000001"));
402 let v = build_zipkin(&span, &test_ctx(), "svc", Some("10.0.0.5"));
403 let s = &v[0];
404 assert_eq!(s["traceId"], "0af7651916cd43dd8448eb211c80319c");
405 assert_eq!(s["id"], "b7ad6b7169203331");
406 assert_eq!(s["parentId"], "0020000000000001");
407 assert_eq!(s["kind"], "SERVER");
408 assert_eq!(s["name"], "POST /api/users");
409 assert_eq!(s["timestamp"], 1_700_000_000_000_000u64);
411 assert_eq!(s["localEndpoint"]["serviceName"], "svc");
412 assert_eq!(s["localEndpoint"]["ipv4"], "10.0.0.5");
413 assert_eq!(s["tags"]["http.method"], "POST");
414 assert_eq!(s["tags"]["http.status_code"], "201");
415 assert_eq!(s["tags"]["http.path"], "/api/users");
416 }
417
418 #[test]
419 fn zipkin_omits_parent_when_root() {
420 let span = a_span(true, None);
421 let v = build_zipkin(&span, &test_ctx(), "svc", None);
422 assert!(v[0].get("parentId").is_none());
423 assert!(v[0]["localEndpoint"].get("ipv4").is_none());
424 }
425
426 #[tokio::test]
427 async fn start_creates_new_trace_and_injects_headers() {
428 let p =
429 plugin(json!({ "phase": "start", "endpoint": "http://z", "sample_ratio": 1 })).unwrap();
430 let out = p.execute(test_ctx(), &HashMap::new()).await.unwrap();
431 let span = load_span(&out.context).expect("span stored");
432 assert!(span.parent_span_id.is_none());
433 assert!(span.sampled);
434 assert_eq!(span.trace_id.len(), 32);
435 assert_eq!(
437 out.context.request.headers["x-b3-traceid"][0],
438 span.trace_id
439 );
440 assert_eq!(out.context.request.headers["x-b3-spanid"][0], span.span_id);
441 assert_eq!(out.context.request.headers["x-b3-sampled"][0], "1");
442 }
443
444 #[tokio::test]
445 async fn start_continues_incoming_b3_single() {
446 let p = plugin(json!({ "phase": "start", "endpoint": "http://z" })).unwrap();
447 let mut ctx = test_ctx();
448 ctx.request
449 .headers
450 .insert("b3".to_string(), vec!["trace1-span1-1".to_string()]);
451 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
452 let span = load_span(&out.context).unwrap();
453 assert_eq!(span.trace_id, "trace1");
454 assert_eq!(span.parent_span_id.as_deref(), Some("span1"));
455 assert!(span.sampled);
456 assert_ne!(span.span_id, "span1");
458 }
459
460 #[tokio::test]
461 async fn start_continues_incoming_b3_multi() {
462 let p = plugin(json!({ "phase": "start", "endpoint": "http://z" })).unwrap();
463 let mut ctx = test_ctx();
464 ctx.request
465 .headers
466 .insert("x-b3-traceid".to_string(), vec!["t9".to_string()]);
467 ctx.request
468 .headers
469 .insert("x-b3-spanid".to_string(), vec!["s9".to_string()]);
470 ctx.request
471 .headers
472 .insert("x-b3-sampled".to_string(), vec!["0".to_string()]);
473 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
474 let span = load_span(&out.context).unwrap();
475 assert_eq!(span.trace_id, "t9");
476 assert_eq!(span.parent_span_id.as_deref(), Some("s9"));
477 assert!(!span.sampled);
478 }
479
480 #[tokio::test]
481 async fn end_passes_through_without_span() {
482 let p = plugin(json!({ "phase": "end", "endpoint": "http://z" })).unwrap();
483 let out = p.execute(test_ctx(), &HashMap::new()).await.unwrap();
484 assert_eq!(out.context.response.status_code, 201);
485 }
486
487 #[tokio::test]
488 async fn end_returns_ok_with_stored_span() {
489 let p = plugin(json!({ "phase": "end", "endpoint": "http://127.0.0.1:1/api/v2/spans" }))
490 .unwrap();
491 let mut ctx = test_ctx();
492 store_span(&mut ctx, &a_span(true, None));
493 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
495 assert_eq!(out.context.request.path, "/api/users");
496 }
497
498 #[tokio::test]
499 async fn end_skips_export_when_unsampled() {
500 let p = plugin(json!({ "phase": "end", "endpoint": "http://z" })).unwrap();
501 let mut ctx = test_ctx();
502 store_span(&mut ctx, &a_span(false, None));
503 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
504 assert_eq!(out.context.response.status_code, 201);
505 }
506}