Pull sampling probability from env var OTEL_TRACES_SAMPLER_ARG in t…

…he default sdk config Signed-off-by: jbtrystram <jbtrystram@redhat.com>
open-telemetry · Feb 18, 2022 · c2ee2a6 · c2ee2a6
1 parent dacd75a
commit c2ee2a6
Showing 1 changed file with 67 additions and 0 deletions.
diff --git a/opentelemetry-sdk/src/trace/config.rs b/opentelemetry-sdk/src/trace/config.rs
@@ -2,6 +2,7 @@
 //!
 //! Configuration represents the global tracing configuration, overrides
 //! can be set for the default OpenTelemetry limits and Sampler.
+use crate::global::{handle_error, Error};
 use crate::trace::{span_limit::SpanLimits, Sampler, ShouldSample};
 use opentelemetry_api::trace::IdGenerator;
 use std::env;
@@ -123,6 +124,72 @@ impl Default for Config {
             config.span_limits.max_links_per_span = max_links_per_span;
         }
 
+        let sampler_arg = env::var("OTEL_TRACES_SAMPLER_ARG").ok();
+        if let Ok(sampler) = env::var("OTEL_TRACES_SAMPLER") {
+            config.sampler = match sampler.as_str() {
+                "always_on" => Box::new(Sampler::AlwaysOn),
+                "always_off" => Box::new(Sampler::AlwaysOff),
+                "traceidratio" => {
+                    let ratio = sampler_arg.and_then(|r| r.parse::<f64>().ok());
+                    if let Some(r) = ratio {
+                        Box::new(Sampler::TraceIdRatioBased(r))
+                    } else {
+                        handle_error(
+                            Error::Other(String::from(
+                                "Missing or invalid OTEL_TRACES_SAMPLER_ARG value. Falling back to default: 1.0"))
+                        );
+                        Box::new(Sampler::TraceIdRatioBased(1.0))
+                    }
+                }
+                "parentbased_always_on" => {
+                    Box::new(Sampler::ParentBased(Box::new(Sampler::AlwaysOn)))
+                }
+                "parentbased_always_off" => {
+                    Box::new(Sampler::ParentBased(Box::new(Sampler::AlwaysOff)))
+                }
+                "parentbased_traceidratio" => {
+                    let ratio = sampler_arg.and_then(|r| r.parse::<f64>().ok());
+                    if let Some(r) = ratio {
+                        Box::new(Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(
+                            r,
+                        ))))
+                    } else {
+                        handle_error(
+                            Error::Other(String::from(
+                            "Missing or invalid OTEL_TRACES_SAMPLER_ARG value. Falling back to default: 1.0"
+                        )));
+                        Box::new(Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(
+                            1.0,
+                        ))))
+                    }
+                }
+                "parentbased_jaeger_remote" => {
+                    handle_error(
+                        Error::Other(String::from(
+                        "Unimplemented parentbased_jaeger_remote sampler. Falling back to default: parentbased_always_on"
+                    )));
+                    Box::new(Sampler::ParentBased(Box::new(Sampler::AlwaysOn)))
+                }
+                "jaeger_remote" => {
+                    handle_error(
+                        Error::Other(String::from("Unimplemented jaeger_remote sampler. Falling back to default: parentbased_always_on")));
+                    Box::new(Sampler::ParentBased(Box::new(Sampler::AlwaysOn)))
+                }
+                "xray" => {
+                    handle_error(
+                        Error::Other(String::from("Unimplemented xray sampler. Falling back to default: parentbased_always_on")));
+                    Box::new(Sampler::ParentBased(Box::new(Sampler::AlwaysOn)))
+                }
+                s => {
+                    handle_error(
+                        Error::Other(format!("Unrecognised OTEL_TRACES_SAMPLER value: {}. Falling back to default: parentbased_always_on",
+                        s
+                    )));
+                    Box::new(Sampler::ParentBased(Box::new(Sampler::AlwaysOn)))
+                }
+            }
+        }
+
         config
     }
 }