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

sz_orm_tracing/
sampling.rs
