scouter-tracing 0.21.1

Helper crate for tracing instrumentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83

use opentelemetry::baggage::BaggageExt;
use opentelemetry::trace::Span;
use opentelemetry::Context;
use opentelemetry::KeyValue;
use opentelemetry_sdk::error::OTelSdkResult;
use opentelemetry_sdk::trace::SpanProcessor;
use opentelemetry_sdk::trace::{BatchConfig as OTelBatchConfig, BatchConfigBuilder};
use pyo3::prelude::*;
use std::time::Duration;

// This is taken from the opentelemetry examples for adding baggage to spans
#[derive(Debug)]
pub(crate) struct EnrichSpanWithBaggageProcessor;

impl SpanProcessor for EnrichSpanWithBaggageProcessor {
    fn force_flush(&self) -> OTelSdkResult {
        Ok(())
    }

    fn shutdown_with_timeout(&self, _timeout: Duration) -> OTelSdkResult {
        Ok(())
    }

    fn on_start(&self, span: &mut opentelemetry_sdk::trace::Span, cx: &Context) {
        for (kk, vv) in cx.baggage().iter() {
            span.set_attribute(KeyValue::new(kk.clone(), vv.0.clone()));
        }
    }

    fn on_end(&self, _span: opentelemetry_sdk::trace::SpanData) {}
}

#[pyclass]
#[derive(PartialEq, Clone, Debug)]
pub struct BatchConfig {
    pub max_queue_size: usize,
    pub scheduled_delay: Duration,
    pub max_export_batch_size: usize,
}

#[pymethods]
impl BatchConfig {
    #[new]
    #[pyo3(signature = (
        max_queue_size=2048,
        scheduled_delay_ms=5000,
        max_export_batch_size=512,
    ))]
    pub fn new(
        max_queue_size: usize,
        scheduled_delay_ms: u64,
        max_export_batch_size: usize,
    ) -> Self {
        BatchConfig {
            max_queue_size,
            scheduled_delay: Duration::from_millis(scheduled_delay_ms),
            max_export_batch_size,
        }
    }
}

impl Default for BatchConfig {
    fn default() -> Self {
        BatchConfig {
            max_queue_size: 2048,
            scheduled_delay: Duration::from_millis(5000),
            max_export_batch_size: 512,
        }
    }
}

impl BatchConfig {
    pub fn to_otlp_config(&self) -> OTelBatchConfig {
        let mut builder = BatchConfigBuilder::default();

        builder = builder
            .with_scheduled_delay(self.scheduled_delay)
            .with_max_queue_size(self.max_queue_size)
            .with_max_export_batch_size(self.max_export_batch_size);

        builder.build()
    }
}