feat(trace): add optional sampling behaviour for routes (#2511)

Co-authored-by: Arun Raj M <jarnura47@gmail.com>

crates/router_env/src/logger/config.rs

@@ -101,6 +101,8 @@ pub struct LogTelemetry {
     pub otel_exporter_otlp_timeout: Option<u64>,
     /// Whether to use xray ID generator, (enable this if you plan to use AWS-XRAY)
     pub use_xray_generator: bool,
+    /// Route Based Tracing
+    pub route_to_trace: Option<Vec<String>>,
 }
 
 /// Telemetry / tracing.

crates/router_env/src/logger/setup.rs

@@ -12,6 +12,7 @@ use opentelemetry::{
         trace::BatchConfig,
         Resource,
     },
+    trace::{TraceContextExt, TraceState},
     KeyValue,
 };
 use opentelemetry_otlp::{TonicExporterBuilder, WithExportConfig};
@@ -132,6 +133,92 @@ fn get_opentelemetry_exporter(config: &config::LogTelemetry) -> TonicExporterBui
     exporter_builder
 }
 
+#[derive(Debug, Clone)]
+enum TraceUrlAssert {
+    Match(String),
+    EndsWith(String),
+}
+
+impl TraceUrlAssert {
+    fn compare_url(&self, url: &str) -> bool {
+        match self {
+            Self::Match(value) => url == value,
+            Self::EndsWith(end) => url.ends_with(end),
+        }
+    }
+}
+
+impl From<String> for TraceUrlAssert {
+    fn from(value: String) -> Self {
+        match value {
+            url if url.starts_with('*') => Self::EndsWith(url.trim_start_matches('*').to_string()),
+            url => Self::Match(url),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+struct TraceAssertion {
+    clauses: Option<Vec<TraceUrlAssert>>,
+    /// default behaviour for tracing if no condition is provided
+    default: bool,
+}
+
+impl TraceAssertion {
+    ///
+    /// Should the provided url be traced
+    ///
+    fn should_trace_url(&self, url: &str) -> bool {
+        match &self.clauses {
+            Some(clauses) => clauses.iter().all(|cur| cur.compare_url(url)),
+            None => self.default,
+        }
+    }
+}
+
+///
+/// Conditional Sampler for providing control on url based tracing
+///
+#[derive(Clone, Debug)]
+struct ConditionalSampler<T: trace::ShouldSample + Clone + 'static>(TraceAssertion, T);
+
+impl<T: trace::ShouldSample + Clone + 'static> trace::ShouldSample for ConditionalSampler<T> {
+    fn should_sample(
+        &self,
+        parent_context: Option<&opentelemetry::Context>,
+        trace_id: opentelemetry::trace::TraceId,
+        name: &str,
+        span_kind: &opentelemetry::trace::SpanKind,
+        attributes: &opentelemetry::trace::OrderMap<opentelemetry::Key, opentelemetry::Value>,
+        links: &[opentelemetry::trace::Link],
+        instrumentation_library: &opentelemetry::InstrumentationLibrary,
+    ) -> opentelemetry::trace::SamplingResult {
+        match attributes
+            .get(&opentelemetry::Key::new("http.route"))
+            .map_or(self.0.default, |inner| {
+                self.0.should_trace_url(&inner.as_str())
+            }) {
+            true => self.1.should_sample(
+                parent_context,
+                trace_id,
+                name,
+                span_kind,
+                attributes,
+                links,
+                instrumentation_library,
+            ),
+            false => opentelemetry::trace::SamplingResult {
+                decision: opentelemetry::trace::SamplingDecision::Drop,
+                attributes: Vec::new(),
+                trace_state: match parent_context {
+                    Some(ctx) => ctx.span().span_context().trace_state().clone(),
+                    None => TraceState::default(),
+                },
+            },
+        }
+    }
+}
+
 fn setup_tracing_pipeline(
     config: &config::LogTelemetry,
     service_name: &str,
@@ -140,9 +227,16 @@ fn setup_tracing_pipeline(
     global::set_text_map_propagator(TraceContextPropagator::new());
 
     let mut trace_config = trace::config()
-        .with_sampler(trace::Sampler::TraceIdRatioBased(
-            config.sampling_rate.unwrap_or(1.0),
-        ))
+        .with_sampler(trace::Sampler::ParentBased(Box::new(ConditionalSampler(
+            TraceAssertion {
+                clauses: config
+                    .route_to_trace
+                    .clone()
+                    .map(|inner| inner.into_iter().map(Into::into).collect()),
+                default: false,
+            },
+            trace::Sampler::TraceIdRatioBased(config.sampling_rate.unwrap_or(1.0)),
+        ))))
         .with_resource(Resource::new(vec![KeyValue::new(
             "service.name",
             service_name.to_owned(),