HDDS-14814. Unify Fragmented Traces for Freon randomkeys Command

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/tracing/LoopSampler.java

@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.hdds.tracing;
+
+import java.util.concurrent.ThreadLocalRandom;
+
+/**
+ * Probability-based span sampler that samples spans independently.
+ * Uses ThreadLocalRandom for probability decisions.
+ */
+public final class LoopSampler {
+  private final double probability;
+
+  public LoopSampler(double ratio) {
+    if (ratio < 0) {
+      throw new IllegalArgumentException("Sampling ratio cannot be negative: " + ratio);
+    }
+    // Cap at 1.0 to prevent logic errors
+    this.probability = Math.min(ratio, 1.0);
+  }
+
+  public boolean shouldSample() {
+    if (probability <= 0) {
+      return false;
+    }
+    if (probability >= 1.0) {
+      return true;
+    }
+    return ThreadLocalRandom.current().nextDouble() < probability;
+  }
+}

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/tracing/SpanSampler.java

@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.hdds.tracing;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Custom Sampler that applies span-level sampling for configured
+ * span names, and delegates to parent-based strategy otherwise.
+ * When a span name is in the configured spanMap, uses LoopSampler for
+ * probabilistic sampling, otherwise follows the parent span's
+ * sampling decision.
+ */
+public final class SpanSampler implements Sampler {
+
+  private final Sampler rootSampler;
+  private final Map<String, LoopSampler> spanMap;
+
+  public SpanSampler(Sampler rootSampler,
+                     Map<String, LoopSampler> spanMap) {
+    this.rootSampler = rootSampler;
+    this.spanMap = spanMap;
+  }
+
+  @Override
+  public SamplingResult shouldSample(
+      Context parentContext,
+      String traceId,
+      String spanName,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+
+    Span parentSpan = Span.fromContext(parentContext);
+
+    // check if we have a valid parent span
+    if (!parentSpan.getSpanContext().isValid()) {
+      // Root span: always delegate to trace-level sampler
+      // This ensures OTEL_TRACES_SAMPLER_ARG is respected
+      return rootSampler.shouldSample(parentContext, traceId, spanName,
+          spanKind, attributes, parentLinks);
+    }
+
+    // Child span: check parent's sampling status first
+    // after the process of sampling trace / parent span then check if it is sampled or not.
+    if (!parentSpan.getSpanContext().isSampled()) {
+      // Parent was not sampled, so this child should not be sampled either
+      // This prevents orphaned spans
+      return SamplingResult.drop();
+    }
+
+    // Parent was sampled, now check if this span has explicit sampling config
+    LoopSampler loopSampler = spanMap.get(spanName);
+    if (loopSampler != null) {
+      boolean sample = loopSampler.shouldSample();
+      return sample ? SamplingResult.recordAndSample() : SamplingResult.drop();
+    }
+
+    // No explicit config for this span, follow parent's sampling decision
+    return SamplingResult.recordAndSample();
+  }
+
+  @Override
+  public String getDescription() {
+    return "SpanSampler(spanMap=" + spanMap.keySet() + ")";
+  }
+}

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/tracing/TracingUtil.java

@@ -33,6 +33,7 @@
 import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import java.lang.reflect.Proxy;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import org.apache.hadoop.hdds.conf.ConfigurationSource;
@@ -52,6 +53,8 @@ public final class TracingUtil {
   private static final String OTEL_EXPORTER_OTLP_ENDPOINT_DEFAULT = "http://localhost:4317";
   private static final String OTEL_TRACES_SAMPLER_ARG = "OTEL_TRACES_SAMPLER_ARG";
   private static final double OTEL_TRACES_SAMPLER_RATIO_DEFAULT = 1.0;
+  private static final String OTEL_SPAN_SAMPLING_ARG = "OTEL_SPAN_SAMPLING_ARG";
+  private static final String OTEL_TRACES_SAMPLER_CONFIG_DEFAULT = "";
 
   private static volatile boolean isInit = false;
   private static Tracer tracer = OpenTelemetry.noop().getTracer("noop");
@@ -88,21 +91,50 @@ private static void initialize(String serviceName) {
       String sampleStrRatio = System.getenv(OTEL_TRACES_SAMPLER_ARG);
       if (sampleStrRatio != null && !sampleStrRatio.isEmpty()) {
         samplerRatio = Double.parseDouble(System.getenv(OTEL_TRACES_SAMPLER_ARG));
+        LOG.info("Sampling Trace Config = '{}'", samplerRatio);
       }
     } catch (NumberFormatException ex) {
-      // ignore and use the default value.
+      // log and use the default value.
+      LOG.warn("Invalid value for {}: '{}'. Falling back to default: {}",
+          OTEL_TRACES_SAMPLER_ARG, System.getenv(OTEL_TRACES_SAMPLER_ARG), OTEL_TRACES_SAMPLER_RATIO_DEFAULT, ex);
     }
 
+    String spanSamplingConfig = OTEL_TRACES_SAMPLER_CONFIG_DEFAULT;
+    try {
+      String spanStrConfig = System.getenv(OTEL_SPAN_SAMPLING_ARG);
+      if (spanStrConfig != null && !spanStrConfig.isEmpty()) {
+        spanSamplingConfig = spanStrConfig;
+      }
+      LOG.info("Sampling Span Config = '{}'", spanSamplingConfig);
+    } catch (Exception ex) {
+      // Log and use the default value.
+      LOG.warn("Failed to process {}. Falling back to default configuration: {}",
+          OTEL_SPAN_SAMPLING_ARG, OTEL_TRACES_SAMPLER_CONFIG_DEFAULT, ex);
+    }
+    // Pass the config to parseSpanSamplingConfig to get spans to be sampled.
+    Map<String, LoopSampler> spanMap = parseSpanSamplingConfig(spanSamplingConfig);
+
     Resource resource = Resource.create(Attributes.of(AttributeKey.stringKey("service.name"), serviceName));
     OtlpGrpcSpanExporter spanExporter = OtlpGrpcSpanExporter.builder()
         .setEndpoint(otelEndPoint)
         .build();
 
     SimpleSpanProcessor spanProcessor = SimpleSpanProcessor.builder(spanExporter).build();
+
+    // Choose sampler based on span sampling config. If it is empty use trace based sampling only.
+    // else use custom SpanSampler.
+    Sampler sampler;
+    if (spanMap.isEmpty()) {
+      sampler = Sampler.traceIdRatioBased(samplerRatio);
+    } else {
+      Sampler rootSampler = Sampler.traceIdRatioBased(samplerRatio);
+      sampler = new SpanSampler(rootSampler, spanMap);
+    }
+
     SdkTracerProvider tracerProvider = SdkTracerProvider.builder()
         .addSpanProcessor(spanProcessor)
         .setResource(resource)
-        .setSampler(Sampler.traceIdRatioBased(samplerRatio))
+        .setSampler(sampler)
         .build();
     OpenTelemetry openTelemetry = OpenTelemetrySdk.builder()
         .setTracerProvider(tracerProvider)
@@ -177,6 +209,43 @@ public static boolean isTracingEnabled(
         ScmConfigKeys.HDDS_TRACING_ENABLED_DEFAULT);
   }
 
+  /** Function to parse span sampling config. The input is in the form <span_name>:<sample_rate>.
+   * The sample rate must be a number between 0 and 1. Any value other than that will LOG an error.
+   * */
+  private static Map<String, LoopSampler> parseSpanSamplingConfig(String configStr) {
+    Map<String, LoopSampler> result = new HashMap<>();
+    if (configStr == null || configStr.isEmpty()) {
+      return Collections.emptyMap();
+    }
+
+    for (String entry : configStr.split(",")) {
+      String trimmed = entry.trim();
+      int colon = trimmed.indexOf(':');
+
+      if (colon <= 0 || colon >= trimmed.length() - 1) {
+        continue;
+      }
+
+      String name = trimmed.substring(0, colon).trim();
+      String val = trimmed.substring(colon + 1).trim();
+
+      try {
+        double rate = Double.parseDouble(val);
+        //if the rate  is less than or equal to zero , no sampling config is taken for that key value pair.
+        if (rate > 0) {
+          // cap it at 1.0 when a number greater than 1 is entered
+          double effectiveRate = Math.min(rate, 1.0);
+          result.put(name, new LoopSampler(effectiveRate));
+        } else {
+          LOG.warn("rate for span '{}' is 0 or less, ignoring sample configuration", name);
+        }
+      } catch (NumberFormatException e) {
+        LOG.error("Invalid rate '{}' for span '{}', ignoring sample configuration", val, name);
+      }
+    }
+    return result;
+  }
+
   /**
    * Execute {@code runnable} inside an activated new span.
    * If a parent span exists in the current context, this becomes a child span.