sz_orm_tracing/
sampling.rs

1//! 采样策略与 Baggage 传播
2//!
3//! 提供分布式追踪的采样决策与 Baggage 上下文传播能力。
4//!
5//! ## 采样策略
6//!
7//! - [`TraceIdRatioSampler`]：基于 trace_id 的确定性采样，保证同一 trace
8//!   在所有节点上做出相同的采样决策。
9//! - [`ParentBasedSampler`]：遵循父 span 的采样标志，适用于局部采样场景。
10//! - [`AlwaysOnSampler`] / [`AlwaysOffSampler`]：全量采样/全量不采样。
11//!
12//! ## Baggage 传播
13//!
14//! Baggage 是 W3C 规范定义的跨进程键值对传播机制，用于在请求链路上传递
15//! 业务上下文（如 user_id、request_id、locale）。
16//!
17//! 格式：`baggage: key1=value1,key2=value2`
18
19use std::collections::HashMap;
20use std::sync::atomic::{AtomicU64, Ordering};
21use std::sync::Arc;
22
23/// 采样决策
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum SamplingDecision {
26    /// 采样：该 trace 会被记录与上报
27    RecordAndSample,
28    /// 不采样：该 trace 不会被记录
29    NotRecord,
30}
31
32impl SamplingDecision {
33    /// 是否采样
34    pub fn is_sampled(&self) -> bool {
35        matches!(self, SamplingDecision::RecordAndSample)
36    }
37
38    /// 转为 W3C traceparent 中的 trace_flags（`01` 表示采样，`00` 表示不采样）
39    pub fn as_trace_flags(&self) -> &'static str {
40        match self {
41            SamplingDecision::RecordAndSample => "01",
42            SamplingDecision::NotRecord => "00",
43        }
44    }
45}
46
47/// 采样器 trait：根据 trace_id 与父 span 状态做出采样决策
48pub trait Sampler: Send + Sync {
49    /// 做出采样决策
50    ///
51    /// # 参数
52    /// - `trace_id`：trace 标识（32 字符 hex）
53    /// - `parent_sampled`：父 span 是否已采样（`None` 表示根 span）
54    fn should_sample(&self, trace_id: &str, parent_sampled: Option<bool>) -> SamplingDecision;
55
56    /// 采样器名称（用于可观测性）
57    fn name(&self) -> &'static str;
58}
59
60/// 全量采样器：始终返回 RecordAndSample
61#[derive(Debug, Clone, Default)]
62pub struct AlwaysOnSampler;
63
64impl AlwaysOnSampler {
65    pub fn new() -> Self {
66        Self
67    }
68}
69
70impl Sampler for AlwaysOnSampler {
71    fn should_sample(&self, _trace_id: &str, _parent_sampled: Option<bool>) -> SamplingDecision {
72        SamplingDecision::RecordAndSample
73    }
74
75    fn name(&self) -> &'static str {
76        "always_on"
77    }
78}
79
80/// 全量不采样器：始终返回 NotRecord
81#[derive(Debug, Clone, Default)]
82pub struct AlwaysOffSampler;
83
84impl AlwaysOffSampler {
85    pub fn new() -> Self {
86        Self
87    }
88}
89
90impl Sampler for AlwaysOffSampler {
91    fn should_sample(&self, _trace_id: &str, _parent_sampled: Option<bool>) -> SamplingDecision {
92        SamplingDecision::NotRecord
93    }
94
95    fn name(&self) -> &'static str {
96        "always_off"
97    }
98}
99
100/// 基于 trace_id 比例的确定性采样器。
101///
102/// 使用 trace_id 的哈希值与阈值比较，保证同一 trace_id 在所有节点上
103/// 做出相同的采样决策（无需共享状态）。
104///
105/// 比例范围为 0.0..=1.0：
106/// - 0.0：不采样任何 trace
107/// - 1.0：采样所有 trace
108/// - 0.5：采样约 50% 的 trace
109pub struct TraceIdRatioSampler {
110    /// 采样比例（0.0..=1.0）
111    ratio: f64,
112    /// 阈值 = ratio * u64::MAX
113    threshold: u64,
114    /// 总采样次数（用于统计）
115    total: AtomicU64,
116    /// 被采样的次数
117    sampled: AtomicU64,
118}
119
120impl TraceIdRatioSampler {
121    /// 创建比例采样器
122    ///
123    /// # Panics
124    /// 当 ratio 不在 [0.0, 1.0] 范围内时 panic。
125    pub fn new(ratio: f64) -> Self {
126        assert!(
127            (0.0..=1.0).contains(&ratio),
128            "sampling ratio must be in [0.0, 1.0], got {ratio}"
129        );
130        Self {
131            ratio,
132            threshold: (ratio * u64::MAX as f64) as u64,
133            total: AtomicU64::new(0),
134            sampled: AtomicU64::new(0),
135        }
136    }
137
138    /// 采样比例
139    pub fn ratio(&self) -> f64 {
140        self.ratio
141    }
142
143    /// 总采样决策次数
144    pub fn total_decisions(&self) -> u64 {
145        self.total.load(Ordering::Relaxed)
146    }
147
148    /// 被采样的次数
149    pub fn sampled_count(&self) -> u64 {
150        self.sampled.load(Ordering::Relaxed)
151    }
152
153    /// 实际采样率（sampled / total），总次数为 0 时返回 0.0
154    pub fn actual_rate(&self) -> f64 {
155        let total = self.total_decisions();
156        if total == 0 {
157            return 0.0;
158        }
159        self.sampled_count() as f64 / total as f64
160    }
161
162    /// 对 trace_id 做确定性哈希，返回 u64 值。
163    ///
164    /// 使用 FNV-1a 64-bit 哈希 + 末级 avalanche 混合（xorshift-mix），
165    /// 保证同一输入始终产生相同输出，且对相近输入（如顺序递增的 hex 字符串）
166    /// 也能均匀分散到 u64 值域，避免比例采样时出现严重偏差。
167    fn hash_trace_id(trace_id: &str) -> u64 {
168        // FNV-1a 64-bit
169        const FNV_OFFSET: u64 = 14695981039346656037;
170        const FNV_PRIME: u64 = 1099511628211;
171
172        let mut hash = FNV_OFFSET;
173        for byte in trace_id.as_bytes() {
174            hash ^= *byte as u64;
175            hash = hash.wrapping_mul(FNV_PRIME);
176        }
177
178        // Avalanche finalizer（基于 xxHash3 的尾混合思想）：
179        // 将高位与低位充分混合，使输出在 u64 域内均匀分布。
180        hash ^= hash >> 33;
181        hash = hash.wrapping_mul(0xff51_afd7_ed55_8ccd);
182        hash ^= hash >> 33;
183        hash = hash.wrapping_mul(0xc4ce_b9fe_1a85_ec53);
184        hash ^= hash >> 33;
185        hash
186    }
187}
188
189impl Sampler for TraceIdRatioSampler {
190    fn should_sample(&self, trace_id: &str, _parent_sampled: Option<bool>) -> SamplingDecision {
191        self.total.fetch_add(1, Ordering::Relaxed);
192        let hash = Self::hash_trace_id(trace_id);
193        if hash <= self.threshold {
194            self.sampled.fetch_add(1, Ordering::Relaxed);
195            SamplingDecision::RecordAndSample
196        } else {
197            SamplingDecision::NotRecord
198        }
199    }
200
201    fn name(&self) -> &'static str {
202        "trace_id_ratio"
203    }
204}
205
206/// 基于父 span 的采样器。
207///
208/// 采样策略：
209/// - 根 span（parent_sampled = None）：使用内部 root sampler 决策
210/// - 有父 span：遵循父 span 的采样标志
211pub struct ParentBasedSampler {
212    /// 根 span 采样器
213    root: Box<dyn Sampler>,
214}
215
216impl ParentBasedSampler {
217    /// 创建基于父 span 的采样器
218    ///
219    /// # 参数
220    /// - `root`：根 span 使用的采样器（通常为 TraceIdRatioSampler）
221    pub fn new(root: Box<dyn Sampler>) -> Self {
222        Self { root }
223    }
224
225    /// 使用 AlwaysOn 作为根采样器的便捷构造
226    pub fn always_on_root() -> Self {
227        Self::new(Box::new(AlwaysOnSampler::new()))
228    }
229
230    /// 使用 TraceIdRatio 作为根采样器的便捷构造
231    pub fn ratio_root(ratio: f64) -> Self {
232        Self::new(Box::new(TraceIdRatioSampler::new(ratio)))
233    }
234}
235
236impl Sampler for ParentBasedSampler {
237    fn should_sample(&self, trace_id: &str, parent_sampled: Option<bool>) -> SamplingDecision {
238        match parent_sampled {
239            None => self.root.should_sample(trace_id, None),
240            Some(sampled) => {
241                if sampled {
242                    SamplingDecision::RecordAndSample
243                } else {
244                    SamplingDecision::NotRecord
245                }
246            }
247        }
248    }
249
250    fn name(&self) -> &'static str {
251        "parent_based"
252    }
253}
254
255// ============================================================================
256// Baggage 传播
257// ============================================================================
258
259/// Baggage：跨进程键值对上下文。
260///
261/// 遵循 W3C Baggage 规范，通过 `baggage` HTTP header 传播。
262/// 格式：`key1=value1,key2=value2`
263#[derive(Debug, Clone, Default)]
264pub struct Baggage {
265    entries: HashMap<String, String>,
266}
267
268impl Baggage {
269    /// 创建空 Baggage
270    pub fn new() -> Self {
271        Self::default()
272    }
273
274    /// 从键值对列表创建
275    pub fn from_pairs(
276        pairs: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>,
277    ) -> Self {
278        let mut baggage = Self::new();
279        for (k, v) in pairs {
280            baggage.set(k, v);
281        }
282        baggage
283    }
284
285    /// 设置键值对（覆盖已有值）
286    pub fn set(&mut self, key: impl Into<String>, value: impl Into<String>) {
287        self.entries.insert(key.into(), value.into());
288    }
289
290    /// 获取值
291    pub fn get(&self, key: &str) -> Option<&str> {
292        self.entries.get(key).map(|s| s.as_str())
293    }
294
295    /// 移除键
296    pub fn remove(&mut self, key: &str) -> Option<String> {
297        self.entries.remove(key)
298    }
299
300    /// 是否为空
301    pub fn is_empty(&self) -> bool {
302        self.entries.is_empty()
303    }
304
305    /// 条目数量
306    pub fn len(&self) -> usize {
307        self.entries.len()
308    }
309
310    /// 获取所有键
311    pub fn keys(&self) -> Vec<&str> {
312        self.entries.keys().map(|s| s.as_str()).collect()
313    }
314
315    /// 序列化为 W3C baggage header 值
316    ///
317    /// 格式：`key1=value1,key2=value2`
318    /// 键值按字母序排列，保证输出确定性。
319    pub fn to_header(&self) -> String {
320        let mut pairs: Vec<(String, String)> = self
321            .entries
322            .iter()
323            .map(|(k, v)| (k.clone(), v.clone()))
324            .collect();
325        pairs.sort_by(|a, b| a.0.cmp(&b.0));
326        pairs
327            .into_iter()
328            .map(|(k, v)| format!("{}={}", k, v))
329            .collect::<Vec<_>>()
330            .join(",")
331    }
332
333    /// 从 W3C baggage header 值解析
334    ///
335    /// 格式：`key1=value1,key2=value2`
336    /// 解析时忽略空白条目和格式不正确的条目（无 `=`、空 key、空 value）。
337    pub fn from_header(header: &str) -> Self {
338        let mut baggage = Self::new();
339        for entry in header.split(',') {
340            let entry = entry.trim();
341            if entry.is_empty() {
342                continue;
343            }
344            if let Some(eq_pos) = entry.find('=') {
345                let key = entry[..eq_pos].trim().to_string();
346                let value = entry[eq_pos + 1..].trim().to_string();
347                // 键和值都必须非空，否则视为格式错误并忽略。
348                if !key.is_empty() && !value.is_empty() {
349                    baggage.entries.insert(key, value);
350                }
351            }
352        }
353        baggage
354    }
355
356    /// 合并另一个 Baggage（后者覆盖前者）
357    pub fn merge(&mut self, other: &Baggage) {
358        for (k, v) in &other.entries {
359            self.entries.insert(k.clone(), v.clone());
360        }
361    }
362
363    /// 清空所有条目
364    pub fn clear(&mut self) {
365        self.entries.clear();
366    }
367}
368
369/// Baggage 传播器：在 HTTP header 与 Baggage 对象之间转换
370pub struct BaggagePropagator;
371
372impl BaggagePropagator {
373    /// 创建传播器
374    pub fn new() -> Self {
375        Self
376    }
377
378    /// 将 Baggage 注入到 header map
379    pub fn inject(&self, baggage: &Baggage, headers: &mut HashMap<String, String>) {
380        let header_value = baggage.to_header();
381        if !header_value.is_empty() {
382            headers.insert("baggage".to_string(), header_value);
383        }
384    }
385
386    /// 从 header map 提取 Baggage
387    pub fn extract(&self, headers: &HashMap<String, String>) -> Baggage {
388        headers
389            .get("baggage")
390            .map(|h| Baggage::from_header(h))
391            .unwrap_or_default()
392    }
393}
394
395impl Default for BaggagePropagator {
396    fn default() -> Self {
397        Self::new()
398    }
399}
400
401// ============================================================================
402// 批量导出
403// ============================================================================
404
405/// 批量 Span 导出器配置
406#[derive(Debug, Clone)]
407pub struct BatchConfig {
408    /// 批次最大大小（达到后立即导出）
409    pub max_batch_size: usize,
410    /// 导出间隔（毫秒）
411    pub export_interval_ms: u64,
412    /// 最大队列长度（超出时丢弃最旧的 span）
413    pub max_queue_size: usize,
414}
415
416impl Default for BatchConfig {
417    fn default() -> Self {
418        Self {
419            max_batch_size: 512,
420            export_interval_ms: 5000,
421            max_queue_size: 2048,
422        }
423    }
424}
425
426/// 批量 Span 导出器（内存模拟）。
427///
428/// 收集 span 到队列，达到批次大小或导出间隔时批量导出。
429/// 实际实现中会通过 OTLP/HTTP 发送到 Collector。
430pub struct BatchSpanExporter {
431    config: BatchConfig,
432    /// 待导出的 span 队列
433    queue: Arc<std::sync::RwLock<Vec<crate::Span>>>,
434    /// 已导出的 span 批次
435    exported: Arc<std::sync::RwLock<Vec<Vec<crate::Span>>>>,
436    /// 被丢弃的 span 数量
437    dropped: AtomicU64,
438}
439
440impl BatchSpanExporter {
441    /// 创建批量导出器
442    pub fn new(config: BatchConfig) -> Self {
443        Self {
444            config,
445            queue: Arc::new(std::sync::RwLock::new(Vec::new())),
446            exported: Arc::new(std::sync::RwLock::new(Vec::new())),
447            dropped: AtomicU64::new(0),
448        }
449    }
450
451    /// 将 span 加入导出队列
452    pub fn enqueue(&self, span: crate::Span) {
453        let mut queue = self.queue.write().unwrap();
454        if queue.len() >= self.config.max_queue_size {
455            // 队列已满，丢弃最旧的 span
456            queue.remove(0);
457            self.dropped.fetch_add(1, Ordering::Relaxed);
458        }
459        queue.push(span);
460    }
461
462    /// 尝试导出一批 span。
463    ///
464    /// 当队列长度 >= max_batch_size 时导出整批，否则不导出。
465    /// 返回导出的 span 数量。
466    pub fn flush_batch(&self) -> usize {
467        let mut queue = self.queue.write().unwrap();
468        if queue.len() < self.config.max_batch_size {
469            return 0;
470        }
471        let batch: Vec<crate::Span> = queue.drain(..self.config.max_batch_size).collect();
472        let count = batch.len();
473        let mut exported = self.exported.write().unwrap();
474        exported.push(batch);
475        count
476    }
477
478    /// 强制导出所有排队中的 span，不论批次大小。
479    pub fn flush_all(&self) -> usize {
480        let mut queue = self.queue.write().unwrap();
481        if queue.is_empty() {
482            return 0;
483        }
484        let batch: Vec<crate::Span> = queue.drain(..).collect();
485        let count = batch.len();
486        let mut exported = self.exported.write().unwrap();
487        exported.push(batch);
488        count
489    }
490
491    /// 获取已导出的批次数
492    pub fn exported_batch_count(&self) -> usize {
493        self.exported.read().unwrap().len()
494    }
495
496    /// 获取已导出的 span 总数
497    pub fn exported_span_count(&self) -> usize {
498        self.exported.read().unwrap().iter().map(|b| b.len()).sum()
499    }
500
501    /// 获取当前队列长度
502    pub fn queue_len(&self) -> usize {
503        self.queue.read().unwrap().len()
504    }
505
506    /// 获取被丢弃的 span 数量
507    pub fn dropped_count(&self) -> u64 {
508        self.dropped.load(Ordering::Relaxed)
509    }
510
511    /// 清空所有状态（队列与已导出记录）
512    pub fn clear(&self) {
513        self.queue.write().unwrap().clear();
514        self.exported.write().unwrap().clear();
515        self.dropped.store(0, Ordering::Relaxed);
516    }
517
518    /// 获取配置引用
519    pub fn config(&self) -> &BatchConfig {
520        &self.config
521    }
522}
523
524#[cfg(test)]
525mod tests {
526    use super::*;
527
528    // ===================== SamplingDecision 测试 =====================
529
530    #[test]
531    fn test_sampling_decision_is_sampled() {
532        assert!(SamplingDecision::RecordAndSample.is_sampled());
533        assert!(!SamplingDecision::NotRecord.is_sampled());
534    }
535
536    #[test]
537    fn test_sampling_decision_as_trace_flags() {
538        assert_eq!(SamplingDecision::RecordAndSample.as_trace_flags(), "01");
539        assert_eq!(SamplingDecision::NotRecord.as_trace_flags(), "00");
540    }
541
542    #[test]
543    fn test_sampling_decision_equality() {
544        assert_eq!(
545            SamplingDecision::RecordAndSample,
546            SamplingDecision::RecordAndSample
547        );
548        assert_ne!(
549            SamplingDecision::RecordAndSample,
550            SamplingDecision::NotRecord
551        );
552    }
553
554    // ===================== AlwaysOnSampler 测试 =====================
555
556    #[test]
557    fn test_always_on_sampler_returns_sampled() {
558        let sampler = AlwaysOnSampler::new();
559        let decision = sampler.should_sample("trace123", None);
560        assert_eq!(decision, SamplingDecision::RecordAndSample);
561    }
562
563    #[test]
564    fn test_always_on_sampler_ignores_parent() {
565        let sampler = AlwaysOnSampler::new();
566        assert!(sampler.should_sample("t", Some(false)).is_sampled());
567        assert!(sampler.should_sample("t", Some(true)).is_sampled());
568    }
569
570    #[test]
571    fn test_always_on_sampler_name() {
572        let sampler = AlwaysOnSampler::new();
573        assert_eq!(sampler.name(), "always_on");
574    }
575
576    // ===================== AlwaysOffSampler 测试 =====================
577
578    #[test]
579    fn test_always_off_sampler_returns_not_sampled() {
580        let sampler = AlwaysOffSampler::new();
581        let decision = sampler.should_sample("trace123", None);
582        assert_eq!(decision, SamplingDecision::NotRecord);
583    }
584
585    #[test]
586    fn test_always_off_sampler_name() {
587        let sampler = AlwaysOffSampler::new();
588        assert_eq!(sampler.name(), "always_off");
589    }
590
591    // ===================== TraceIdRatioSampler 测试 =====================
592
593    #[test]
594    fn test_trace_id_ratio_sampler_full_sampling() {
595        let sampler = TraceIdRatioSampler::new(1.0);
596        // ratio=1.0 时所有 trace 都应被采样
597        for i in 0..100 {
598            let trace_id = format!("{:032x}", i);
599            assert!(sampler.should_sample(&trace_id, None).is_sampled());
600        }
601        assert_eq!(sampler.sampled_count(), 100);
602        assert_eq!(sampler.total_decisions(), 100);
603    }
604
605    #[test]
606    fn test_trace_id_ratio_sampler_zero_sampling() {
607        let sampler = TraceIdRatioSampler::new(0.0);
608        // ratio=0.0 时不应采样任何 trace
609        for i in 0..100 {
610            let trace_id = format!("{:032x}", i);
611            assert!(!sampler.should_sample(&trace_id, None).is_sampled());
612        }
613        assert_eq!(sampler.sampled_count(), 0);
614    }
615
616    #[test]
617    fn test_trace_id_ratio_sampler_deterministic() {
618        let sampler = TraceIdRatioSampler::new(0.5);
619        // 同一 trace_id 多次采样应得到相同结果
620        let trace_id = "abcdef0123456789abcdef0123456789";
621        let d1 = sampler.should_sample(trace_id, None);
622        let d2 = sampler.should_sample(trace_id, None);
623        let d3 = sampler.should_sample(trace_id, None);
624        assert_eq!(d1, d2);
625        assert_eq!(d2, d3);
626    }
627
628    #[test]
629    fn test_trace_id_ratio_sampler_half_ratio_approximate() {
630        let sampler = TraceIdRatioSampler::new(0.5);
631        // 大量样本下采样率应接近 50%
632        for i in 0..10000 {
633            let trace_id = format!("{:032x}", i);
634            sampler.should_sample(&trace_id, None);
635        }
636        let rate = sampler.actual_rate();
637        assert!((rate - 0.5).abs() < 0.05, "expected ~0.5, got {rate}");
638    }
639
640    #[test]
641    fn test_trace_id_ratio_sampler_stats() {
642        let sampler = TraceIdRatioSampler::new(0.3);
643        for i in 0..1000 {
644            let trace_id = format!("{:032x}", i);
645            sampler.should_sample(&trace_id, None);
646        }
647        assert_eq!(sampler.total_decisions(), 1000);
648        assert!(sampler.sampled_count() > 0);
649        let rate = sampler.actual_rate();
650        assert!((rate - 0.3).abs() < 0.05, "expected ~0.3, got {rate}");
651    }
652
653    #[test]
654    fn test_trace_id_ratio_sampler_name() {
655        let sampler = TraceIdRatioSampler::new(1.0);
656        assert_eq!(sampler.name(), "trace_id_ratio");
657    }
658
659    #[test]
660    fn test_trace_id_ratio_sampler_ratio_accessor() {
661        let sampler = TraceIdRatioSampler::new(0.75);
662        assert!((sampler.ratio() - 0.75).abs() < 1e-9);
663    }
664
665    #[test]
666    #[should_panic(expected = "sampling ratio must be in [0.0, 1.0]")]
667    fn test_trace_id_ratio_sampler_invalid_ratio_high() {
668        TraceIdRatioSampler::new(1.5);
669    }
670
671    #[test]
672    #[should_panic(expected = "sampling ratio must be in [0.0, 1.0]")]
673    fn test_trace_id_ratio_sampler_invalid_ratio_negative() {
674        TraceIdRatioSampler::new(-0.1);
675    }
676
677    // ===================== ParentBasedSampler 测试 =====================
678
679    #[test]
680    fn test_parent_based_root_uses_inner_sampler() {
681        let sampler = ParentBasedSampler::ratio_root(1.0);
682        // 根 span -> 使用 TraceIdRatio(1.0) -> 采样
683        assert!(sampler.should_sample("trace1", None).is_sampled());
684    }
685
686    #[test]
687    fn test_parent_based_follows_sampled_parent() {
688        let sampler = ParentBasedSampler::always_on_root();
689        // 父 span 已采样 -> 子 span 也采样
690        assert!(sampler.should_sample("t", Some(true)).is_sampled());
691    }
692
693    #[test]
694    fn test_parent_based_follows_unsampled_parent() {
695        let sampler = ParentBasedSampler::always_on_root();
696        // 父 span 未采样 -> 子 span 也不采样（即使 root 是 AlwaysOn）
697        assert!(!sampler.should_sample("t", Some(false)).is_sampled());
698    }
699
700    #[test]
701    fn test_parent_based_name() {
702        let sampler = ParentBasedSampler::always_on_root();
703        assert_eq!(sampler.name(), "parent_based");
704    }
705
706    // ===================== Baggage 测试 =====================
707
708    #[test]
709    fn test_baggage_new_empty() {
710        let b = Baggage::new();
711        assert!(b.is_empty());
712        assert_eq!(b.len(), 0);
713    }
714
715    #[test]
716    fn test_baggage_set_and_get() {
717        let mut b = Baggage::new();
718        b.set("user_id", "12345");
719        assert_eq!(b.get("user_id"), Some("12345"));
720        assert_eq!(b.get("missing"), None);
721    }
722
723    #[test]
724    fn test_baggage_set_overwrites() {
725        let mut b = Baggage::new();
726        b.set("key", "v1");
727        b.set("key", "v2");
728        assert_eq!(b.get("key"), Some("v2"));
729        assert_eq!(b.len(), 1);
730    }
731
732    #[test]
733    fn test_baggage_remove() {
734        let mut b = Baggage::new();
735        b.set("key", "value");
736        assert_eq!(b.remove("key"), Some("value".to_string()));
737        assert!(b.is_empty());
738        assert_eq!(b.remove("key"), None);
739    }
740
741    #[test]
742    fn test_baggage_from_pairs() {
743        let b = Baggage::from_pairs([("a", "1"), ("b", "2")]);
744        assert_eq!(b.len(), 2);
745        assert_eq!(b.get("a"), Some("1"));
746        assert_eq!(b.get("b"), Some("2"));
747    }
748
749    #[test]
750    fn test_baggage_to_header_single() {
751        let mut b = Baggage::new();
752        b.set("key", "value");
753        assert_eq!(b.to_header(), "key=value");
754    }
755
756    #[test]
757    fn test_baggage_to_header_multiple_sorted() {
758        let mut b = Baggage::new();
759        b.set("zebra", "1");
760        b.set("alpha", "2");
761        // 输出按 key 字母序排列
762        assert_eq!(b.to_header(), "alpha=2,zebra=1");
763    }
764
765    #[test]
766    fn test_baggage_to_header_empty() {
767        let b = Baggage::new();
768        assert_eq!(b.to_header(), "");
769    }
770
771    #[test]
772    fn test_baggage_from_header_single() {
773        let b = Baggage::from_header("key=value");
774        assert_eq!(b.get("key"), Some("value"));
775    }
776
777    #[test]
778    fn test_baggage_from_header_multiple() {
779        let b = Baggage::from_header("a=1,b=2,c=3");
780        assert_eq!(b.len(), 3);
781        assert_eq!(b.get("a"), Some("1"));
782        assert_eq!(b.get("b"), Some("2"));
783        assert_eq!(b.get("c"), Some("3"));
784    }
785
786    #[test]
787    fn test_baggage_from_header_with_spaces() {
788        let b = Baggage::from_header(" key = value1 , b = value2 ");
789        assert_eq!(b.get("key"), Some("value1"));
790        assert_eq!(b.get("b"), Some("value2"));
791    }
792
793    #[test]
794    fn test_baggage_from_header_empty() {
795        let b = Baggage::from_header("");
796        assert!(b.is_empty());
797    }
798
799    #[test]
800    fn test_baggage_from_header_ignores_malformed() {
801        let b = Baggage::from_header("valid=1,invalid,=nokey,novalue=,good=2");
802        assert_eq!(b.get("valid"), Some("1"));
803        assert_eq!(b.get("good"), Some("2"));
804        assert_eq!(b.len(), 2);
805    }
806
807    #[test]
808    fn test_baggage_roundtrip() {
809        let mut original = Baggage::new();
810        original.set("user_id", "12345");
811        original.set("request_id", "abc");
812        original.set("locale", "zh-CN");
813
814        let header = original.to_header();
815        let parsed = Baggage::from_header(&header);
816
817        assert_eq!(parsed.len(), original.len());
818        for key in original.keys() {
819            assert_eq!(parsed.get(key), original.get(key));
820        }
821    }
822
823    #[test]
824    fn test_baggage_merge() {
825        let mut b1 = Baggage::new();
826        b1.set("a", "1");
827        b1.set("b", "2");
828
829        let mut b2 = Baggage::new();
830        b2.set("b", "override");
831        b2.set("c", "3");
832
833        b1.merge(&b2);
834        assert_eq!(b1.get("a"), Some("1"));
835        assert_eq!(b1.get("b"), Some("override"));
836        assert_eq!(b1.get("c"), Some("3"));
837    }
838
839    #[test]
840    fn test_baggage_clear() {
841        let mut b = Baggage::new();
842        b.set("a", "1");
843        b.set("b", "2");
844        b.clear();
845        assert!(b.is_empty());
846    }
847
848    #[test]
849    fn test_baggage_keys() {
850        let mut b = Baggage::new();
851        b.set("x", "1");
852        b.set("y", "2");
853        let mut keys = b.keys();
854        keys.sort();
855        assert_eq!(keys, vec!["x", "y"]);
856    }
857
858    // ===================== BaggagePropagator 测试 =====================
859
860    #[test]
861    fn test_propagator_inject_and_extract() {
862        let propagator = BaggagePropagator::new();
863        let mut baggage = Baggage::new();
864        baggage.set("user_id", "123");
865        baggage.set("locale", "en");
866
867        let mut headers = HashMap::new();
868        propagator.inject(&baggage, &mut headers);
869
870        assert!(headers.contains_key("baggage"));
871
872        let extracted = propagator.extract(&headers);
873        assert_eq!(extracted.get("user_id"), Some("123"));
874        assert_eq!(extracted.get("locale"), Some("en"));
875    }
876
877    #[test]
878    fn test_propagator_extract_empty_headers() {
879        let propagator = BaggagePropagator::new();
880        let headers = HashMap::new();
881        let baggage = propagator.extract(&headers);
882        assert!(baggage.is_empty());
883    }
884
885    #[test]
886    fn test_propagator_inject_empty_baggage() {
887        let propagator = BaggagePropagator::new();
888        let baggage = Baggage::new();
889        let mut headers = HashMap::new();
890        propagator.inject(&baggage, &mut headers);
891        // 空 baggage 不应注入 header
892        assert!(!headers.contains_key("baggage"));
893    }
894
895    #[test]
896    fn test_propagator_roundtrip_multiple_entries() {
897        let propagator = BaggagePropagator::new();
898        let mut original = Baggage::new();
899        original.set("a", "1");
900        original.set("b", "2");
901        original.set("c", "3");
902
903        let mut headers = HashMap::new();
904        propagator.inject(&original, &mut headers);
905        let extracted = propagator.extract(&headers);
906
907        assert_eq!(extracted.len(), 3);
908        assert_eq!(extracted.get("a"), Some("1"));
909        assert_eq!(extracted.get("b"), Some("2"));
910        assert_eq!(extracted.get("c"), Some("3"));
911    }
912
913    // ===================== BatchSpanExporter 测试 =====================
914
915    fn make_span(id: &str) -> crate::Span {
916        crate::Span::new("trace", id, "operation")
917    }
918
919    #[test]
920    fn test_batch_exporter_new_empty() {
921        let exporter = BatchSpanExporter::new(BatchConfig::default());
922        assert_eq!(exporter.queue_len(), 0);
923        assert_eq!(exporter.exported_batch_count(), 0);
924        assert_eq!(exporter.exported_span_count(), 0);
925        assert_eq!(exporter.dropped_count(), 0);
926    }
927
928    #[test]
929    fn test_batch_exporter_enqueue() {
930        let exporter = BatchSpanExporter::new(BatchConfig::default());
931        exporter.enqueue(make_span("span1"));
932        exporter.enqueue(make_span("span2"));
933        assert_eq!(exporter.queue_len(), 2);
934    }
935
936    #[test]
937    fn test_batch_exporter_flush_batch_below_threshold() {
938        let config = BatchConfig {
939            max_batch_size: 10,
940            ..Default::default()
941        };
942        let exporter = BatchSpanExporter::new(config);
943        exporter.enqueue(make_span("s1"));
944        exporter.enqueue(make_span("s2"));
945
946        // 队列未达到批次大小，不应导出
947        let exported = exporter.flush_batch();
948        assert_eq!(exported, 0);
949        assert_eq!(exporter.queue_len(), 2);
950    }
951
952    #[test]
953    fn test_batch_exporter_flush_batch_at_threshold() {
954        let config = BatchConfig {
955            max_batch_size: 3,
956            ..Default::default()
957        };
958        let exporter = BatchSpanExporter::new(config);
959        exporter.enqueue(make_span("s1"));
960        exporter.enqueue(make_span("s2"));
961        exporter.enqueue(make_span("s3"));
962
963        let exported = exporter.flush_batch();
964        assert_eq!(exported, 3);
965        assert_eq!(exporter.queue_len(), 0);
966        assert_eq!(exporter.exported_batch_count(), 1);
967        assert_eq!(exporter.exported_span_count(), 3);
968    }
969
970    #[test]
971    fn test_batch_exporter_flush_all() {
972        let exporter = BatchSpanExporter::new(BatchConfig::default());
973        exporter.enqueue(make_span("s1"));
974        exporter.enqueue(make_span("s2"));
975
976        let exported = exporter.flush_all();
977        assert_eq!(exported, 2);
978        assert_eq!(exporter.queue_len(), 0);
979        assert_eq!(exporter.exported_batch_count(), 1);
980    }
981
982    #[test]
983    fn test_batch_exporter_flush_all_empty() {
984        let exporter = BatchSpanExporter::new(BatchConfig::default());
985        let exported = exporter.flush_all();
986        assert_eq!(exported, 0);
987    }
988
989    #[test]
990    fn test_batch_exporter_drops_when_queue_full() {
991        let config = BatchConfig {
992            max_batch_size: 100, // 大批次，不触发自动导出
993            max_queue_size: 3,
994            ..Default::default()
995        };
996        let exporter = BatchSpanExporter::new(config);
997        exporter.enqueue(make_span("s1"));
998        exporter.enqueue(make_span("s2"));
999        exporter.enqueue(make_span("s3"));
1000        // 队列已满，第 4 个应导致丢弃最旧的
1001        exporter.enqueue(make_span("s4"));
1002
1003        assert_eq!(exporter.queue_len(), 3);
1004        assert_eq!(exporter.dropped_count(), 1);
1005    }
1006
1007    #[test]
1008    fn test_batch_exporter_multiple_batches() {
1009        let config = BatchConfig {
1010            max_batch_size: 2,
1011            ..Default::default()
1012        };
1013        let exporter = BatchSpanExporter::new(config);
1014
1015        for i in 0..6 {
1016            exporter.enqueue(make_span(&format!("s{i}")));
1017        }
1018
1019        // 6 个 span，batch_size=2，应导出 3 批
1020        let mut total_exported = 0;
1021        loop {
1022            let n = exporter.flush_batch();
1023            if n == 0 {
1024                break;
1025            }
1026            total_exported += n;
1027        }
1028        assert_eq!(total_exported, 6);
1029        assert_eq!(exporter.exported_batch_count(), 3);
1030    }
1031
1032    #[test]
1033    fn test_batch_exporter_clear() {
1034        let exporter = BatchSpanExporter::new(BatchConfig::default());
1035        exporter.enqueue(make_span("s1"));
1036        exporter.flush_all();
1037
1038        exporter.clear();
1039        assert_eq!(exporter.queue_len(), 0);
1040        assert_eq!(exporter.exported_batch_count(), 0);
1041        assert_eq!(exporter.dropped_count(), 0);
1042    }
1043
1044    #[test]
1045    fn test_batch_config_default() {
1046        let config = BatchConfig::default();
1047        assert_eq!(config.max_batch_size, 512);
1048        assert_eq!(config.export_interval_ms, 5000);
1049        assert_eq!(config.max_queue_size, 2048);
1050    }
1051
1052    #[test]
1053    fn test_batch_exporter_config_accessor() {
1054        let config = BatchConfig {
1055            max_batch_size: 42,
1056            ..Default::default()
1057        };
1058        let exporter = BatchSpanExporter::new(config);
1059        assert_eq!(exporter.config().max_batch_size, 42);
1060    }
1061}
sz_orm_tracing/sampling.rs

sz_orm_tracing/
sampling.rs
