scouter_tracing/exporter/
processor.rs

1use opentelemetry::baggage::BaggageExt;
2use opentelemetry::trace::Span;
3use opentelemetry::Context;
4use opentelemetry::KeyValue;
5use opentelemetry_sdk::error::OTelSdkResult;
6use opentelemetry_sdk::trace::SpanProcessor;
7use opentelemetry_sdk::trace::{BatchConfig as OTelBatchConfig, BatchConfigBuilder};
8use pyo3::prelude::*;
9use std::time::Duration;
10
11// This is taken from the opentelemetry examples for adding baggage to spans
12#[derive(Debug)]
13pub(crate) struct EnrichSpanWithBaggageProcessor;
14
15impl SpanProcessor for EnrichSpanWithBaggageProcessor {
16    fn force_flush(&self) -> OTelSdkResult {
17        Ok(())
18    }
19
20    fn shutdown_with_timeout(&self, _timeout: Duration) -> OTelSdkResult {
21        Ok(())
22    }
23
24    fn on_start(&self, span: &mut opentelemetry_sdk::trace::Span, cx: &Context) {
25        for (kk, vv) in cx.baggage().iter() {
26            span.set_attribute(KeyValue::new(kk.clone(), vv.0.clone()));
27        }
28    }
29
30    fn on_end(&self, _span: opentelemetry_sdk::trace::SpanData) {}
31}
32
33#[pyclass]
34#[derive(PartialEq, Clone, Debug)]
35pub struct BatchConfig {
36    pub max_queue_size: usize,
37    pub scheduled_delay: Duration,
38    pub max_export_batch_size: usize,
39}
40
41#[pymethods]
42impl BatchConfig {
43    #[new]
44    #[pyo3(signature = (
45        max_queue_size=2048,
46        scheduled_delay_ms=5000,
47        max_export_batch_size=512,
48    ))]
49    pub fn new(
50        max_queue_size: usize,
51        scheduled_delay_ms: u64,
52        max_export_batch_size: usize,
53    ) -> Self {
54        BatchConfig {
55            max_queue_size,
56            scheduled_delay: Duration::from_millis(scheduled_delay_ms),
57            max_export_batch_size,
58        }
59    }
60}
61
62impl Default for BatchConfig {
63    fn default() -> Self {
64        BatchConfig {
65            max_queue_size: 2048,
66            scheduled_delay: Duration::from_millis(5000),
67            max_export_batch_size: 512,
68        }
69    }
70}
71
72impl BatchConfig {
73    pub fn to_otlp_config(&self) -> OTelBatchConfig {
74        let mut builder = BatchConfigBuilder::default();
75
76        builder = builder
77            .with_scheduled_delay(self.scheduled_delay)
78            .with_max_queue_size(self.max_queue_size)
79            .with_max_export_batch_size(self.max_export_batch_size);
80
81        builder.build()
82    }
83}
scouter_tracing/exporter/processor.rs

scouter_tracing/exporter/
processor.rs
