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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105

pub mod grpc;
pub mod http;
pub mod noop;
pub mod processor;
pub mod scouter;
pub mod stdout;
pub mod testing;
pub mod traits;

use crate::error::TraceError;
use crate::exporter::noop::NoopSpanExporter;
use crate::exporter::processor::BatchConfig;
use crate::exporter::scouter::ScouterSpanExporter;
use crate::exporter::traits::SpanExporterBuilder;
use opentelemetry_sdk::Resource;
use pyo3::prelude::*;
use scouter_state::app_state;
use tracing::debug;

pub use grpc::GrpcSpanExporter;
pub use http::HttpSpanExporter;
pub use stdout::StdoutSpanExporter;
pub use testing::TestSpanExporter;

#[derive(PartialEq)]
pub enum ExporterType {
    Http,
    Stdout,
    Testing,
    Noop,
    Grpc,
}

// Enum for handling different span exporter types
#[derive(Debug)]
pub enum SpanExporterNum {
    Http(HttpSpanExporter),
    Stdout(StdoutSpanExporter),
    Testing(TestSpanExporter),
    Noop(NoopSpanExporter),
    Grpc(GrpcSpanExporter),
}

impl SpanExporterNum {
    pub fn from_pyobject(obj: &Bound<'_, PyAny>) -> Result<Self, TraceError> {
        if obj.is_instance_of::<HttpSpanExporter>() {
            let exporter = obj.extract::<HttpSpanExporter>()?;
            Ok(SpanExporterNum::Http(exporter))
        } else if obj.is_instance_of::<GrpcSpanExporter>() {
            let exporter = obj.extract::<GrpcSpanExporter>()?;
            Ok(SpanExporterNum::Grpc(exporter))
        } else if obj.is_instance_of::<StdoutSpanExporter>() {
            let exporter = obj.extract::<StdoutSpanExporter>()?;
            Ok(SpanExporterNum::Stdout(exporter))
        } else if obj.is_instance_of::<TestSpanExporter>() {
            let exporter = obj.extract::<TestSpanExporter>()?;
            Ok(SpanExporterNum::Testing(exporter))
        } else {
            debug!("Using NoopSpanExporter as default");
            Ok(SpanExporterNum::Noop(NoopSpanExporter::default()))
        }
    }

    pub(crate) fn set_sample_ratio(&mut self, sample_ratio: Option<f64>) {
        match self {
            SpanExporterNum::Http(builder) => builder.set_sample_ratio(sample_ratio),
            SpanExporterNum::Stdout(builder) => builder.set_sample_ratio(sample_ratio),
            SpanExporterNum::Testing(builder) => builder.set_sample_ratio(sample_ratio),
            SpanExporterNum::Noop(builder) => builder.set_sample_ratio(sample_ratio),
            SpanExporterNum::Grpc(builder) => builder.set_sample_ratio(sample_ratio),
        }
    }

    pub fn build_provider(
        &self,
        resource: Resource,
        scouter_exporter: ScouterSpanExporter,
        batch_config: Option<BatchConfig>,
    ) -> Result<opentelemetry_sdk::trace::SdkTracerProvider, TraceError> {
        match self {
            SpanExporterNum::Http(builder) => {
                builder.build_provider(resource, scouter_exporter, batch_config)
            }
            SpanExporterNum::Stdout(builder) => {
                builder.build_provider(resource, scouter_exporter, batch_config)
            }
            SpanExporterNum::Testing(builder) => {
                builder.build_provider(resource, scouter_exporter, batch_config)
            }
            SpanExporterNum::Noop(builder) => {
                builder.build_provider(resource, scouter_exporter, batch_config)
            }
            // tonic requires a tokio runtime to start the background channel
            SpanExporterNum::Grpc(builder) => app_state().block_on(async {
                builder.build_provider(resource, scouter_exporter, batch_config)
            }),
        }
    }
}

impl Default for SpanExporterNum {
    fn default() -> Self {
        SpanExporterNum::Noop(NoopSpanExporter::default())
    }
}