sampling scaffolding (#830)

* sampling scaffolding * sampling scaffolding * how-to * sampling scaffolding
grafana · Dec 3, 2024 · 03f29b0 · 03f29b0
1 parent 7925bae
commit 03f29b0
Show file tree

Hide file tree

Showing 13 changed files with 309 additions and 21 deletions.
diff --git a/HACKATHON.md b/HACKATHON.md
@@ -0,0 +1,19 @@
+## Building
+
+1. Check out https://github.com/zeitlinger/opentelemetry-java-instrumentation/tree/http-client-response-consumer
+2. Run `./gradlew publishToMavenLocal` in the root directory of the project
+3. Go to the root directory of this project
+4. Apply the patch `git apply opentelemetry-java-instrumentation.patch`
+5. Run `./gradlew clean assemble`
+6. Copy `agent/build/libs/grafana-opentelemetry-java.jar` to wherever you want to use the agent jar,
+    e.g. `cp agent/build/libs/grafana-opentelemetry-java.jar ~/source/docker-otel-lgtm/examples/java`
+
+## Running
+
+1. Either [build the grafana java distro](#building) or download the
+   [latest version](https://drive.google.com/file/d/1koGEMrqZZO6PEcD-hqYJvXYXLy_iCVM3/view?usp=sharing)
+2. Check out https://github.com/grafana/docker-otel-lgtm/tree/hackathon-et-phone-home
+3. The `grafana-opentelemetry-java.jar` from the previous step should be in the `examples/java` directory 
+4. `run-lgtm.sh`
+5. `run-example.sh`
+6. `generate-traffic.sh`
diff --git a/build-local.patch b/build-local.patch
@@ -0,0 +1,50 @@
+Index: build.gradle
+IDEA additional info:
+Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
+<+>UTF-8
+===================================================================
+diff --git a/build.gradle b/build.gradle
+--- a/build.gradle	(revision 0af818d3f2b8a15b97351ef5464e3da5a527f157)
++++ b/build.gradle	(date 1733222424998)
+@@ -1,11 +1,14 @@
++import java.util.regex.Pattern
++
+ group 'com.grafana'
+ version '10.0.0-SNAPSHOT' // the version of the actual release is set during the release process in build-release.sh
+
+ buildscript {
+   ext {
+-    otelInstrumentationVersion = "2.10.0"
++    otelInstrumentationVersion = "2.11.0-SNAPSHOT"
+   }
+   repositories {
++    mavenLocal()
+     maven {
+       url "https://plugins.gradle.org/m2/"
+     }
+@@ -13,7 +16,7 @@
+   dependencies {
+     classpath "com.diffplug.spotless:spotless-plugin-gradle:6.25.0"
+     classpath "com.gradleup.shadow:shadow-gradle-plugin:8.3.5"
+-    classpath "io.opentelemetry.instrumentation:gradle-plugins:${otelInstrumentationVersion}-alpha"
++    classpath "io.opentelemetry.instrumentation:gradle-plugins:2.10.0-alpha"
+   }
+ }
+
+@@ -26,7 +29,7 @@
+   ext {
+     versions = [
+       opentelemetryJavaagent     : otelInstrumentationVersion,
+-      opentelemetryJavaagentAlpha: "${otelInstrumentationVersion}-alpha",
++      opentelemetryJavaagentAlpha: otelInstrumentationVersion.replaceFirst(Pattern.compile("(-SNAPSHOT)?\$"), "-alpha\$1"),
+
+       bytebuddy                  : "1.15.10",
+       autoservice                : "1.1.1",
+@@ -45,6 +48,7 @@
+   }
+
+   repositories {
++    mavenLocal()
+     mavenCentral()
+   }
+
diff --git a/custom/src/main/java/com/grafana/extensions/GrafanaAutoConfigCustomizerProvider.java b/custom/src/main/java/com/grafana/extensions/GrafanaAutoConfigCustomizerProvider.java
@@ -8,6 +8,9 @@
 import com.grafana.extensions.filter.MetricsCustomizer;
 import com.grafana.extensions.instrumentations.TestedInstrumentationsCustomizer;
 import com.grafana.extensions.resources.ResourceCustomizer;
+import com.grafana.extensions.sampler.DeferredSampler;
+import com.grafana.extensions.sampler.SamplingExporter;
+import com.grafana.extensions.sampler.SamplingSpanProcessor;
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer;
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizerProvider;
 import java.util.HashMap;
@@ -21,6 +24,9 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) {
         .addPropertiesSupplier(GrafanaAutoConfigCustomizerProvider::getDefaultProperties)
         .addPropertiesCustomizer(TestedInstrumentationsCustomizer::customizeProperties)
         .addMeterProviderCustomizer(MetricsCustomizer::configure)
+        .addTracerProviderCustomizer(SamplingSpanProcessor::configure)
+        .addSpanExporterCustomizer(SamplingExporter::configure)
+        .addSamplerCustomizer(DeferredSampler::configure)
         .addResourceCustomizer(ResourceCustomizer::truncate);
   }
 

diff --git a/custom/src/main/java/com/grafana/extensions/sampler/DeferredSampler.java b/custom/src/main/java/com/grafana/extensions/sampler/DeferredSampler.java
@@ -0,0 +1,50 @@
+/*
+ * Copyright Grafana Labs
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.grafana.extensions.sampler;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+import java.util.List;
+
+public class DeferredSampler implements Sampler {
+
+  private static final Sampler FROM_PARENT = Sampler.parentBased(Sampler.alwaysOff());
+
+  public static Sampler configure(Sampler configured, ConfigProperties configProperties) {
+    return new DeferredSampler();
+  }
+
+  @Override
+  public String getDescription() {
+    return "DeferredSampler";
+  }
+
+  @Override
+  public SamplingResult shouldSample(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+    SamplingDecision parentDecision =
+        FROM_PARENT
+            .shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks)
+            .getDecision();
+    if (SamplingDecision.RECORD_AND_SAMPLE.equals(parentDecision)) {
+      DynamicSampler.setSampled(traceId);
+    }
+
+    // always return true - because a child span might be sampled even if the parent is not
+    return SamplingResult.recordAndSample();
+  }
+}
diff --git a/custom/src/main/java/com/grafana/extensions/sampler/DynamicSampler.java b/custom/src/main/java/com/grafana/extensions/sampler/DynamicSampler.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright Grafana Labs
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.grafana.extensions.sampler;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import java.util.Collections;
+import java.util.Set;
+import java.util.concurrent.ConcurrentSkipListSet;
+
+public class DynamicSampler {
+  private static final Set<String> sampledTraces = new ConcurrentSkipListSet<>();
+
+  public DynamicSampler(ConfigProperties properties) {
+    // read properties and configure dynamic sampling
+  }
+
+  public static void setSampled(String traceId) {
+    sampledTraces.add(traceId);
+  }
+
+  static boolean isSampled(String traceId) {
+    return sampledTraces.contains(traceId);
+  }
+
+  public static boolean evaluateSampled(ReadableSpan span) {
+    String traceId = span.getSpanContext().getTraceId();
+    if (sampledTraces.contains(traceId)) {
+      return true;
+    }
+    if (shouldSample(span)) {
+      setSampled(traceId);
+      return true;
+    }
+    return false;
+  }
+
+  // public visible for testing
+  public static void clear() {
+    sampledTraces.clear();
+  }
+
+  // visible for testing
+  public static Set<String> getSampledTraces() {
+    return Collections.unmodifiableSet(sampledTraces);
+  }
+
+  static boolean shouldSample(ReadableSpan span) {
+    // add dynamic sampling logic here
+    // dummy implementation for testing for now
+    return Boolean.TRUE.equals(span.getAttributes().get(AttributeKey.booleanKey("sampled")));
+  }
+}
diff --git a/custom/src/main/java/com/grafana/extensions/sampler/SamplingExporter.java b/custom/src/main/java/com/grafana/extensions/sampler/SamplingExporter.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright Grafana Labs
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.grafana.extensions.sampler;
+
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.export.SpanExporter;
+import java.util.ArrayList;
+import java.util.Collection;
+
+public class SamplingExporter implements SpanExporter {
+
+  private final SpanExporter delegate;
+
+  public SamplingExporter(SpanExporter delegate) {
+    this.delegate = delegate;
+  }
+
+  public static SpanExporter configure(SpanExporter delegate, ConfigProperties properties) {
+    return new SamplingExporter(delegate);
+  }
+
+  @Override
+  public CompletableResultCode export(Collection<SpanData> collection) {
+    ArrayList<SpanData> export = new ArrayList<>();
+    for (SpanData data : collection) {
+      if (DynamicSampler.isSampled(data.getSpanContext().getTraceId())) {
+        export.add(data);
+      }
+    }
+    DynamicSampler.clear();
+    return delegate.export(export);
+  }
+
+  @Override
+  public CompletableResultCode flush() {
+    return delegate.flush();
+  }
+
+  @Override
+  public CompletableResultCode shutdown() {
+    return delegate.shutdown();
+  }
+
+  @Override
+  public void close() {
+    delegate.close();
+  }
+}
diff --git a/custom/src/main/java/com/grafana/extensions/sampler/SamplingSpanProcessor.java b/custom/src/main/java/com/grafana/extensions/sampler/SamplingSpanProcessor.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright Grafana Labs
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.grafana.extensions.sampler;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.SdkTracerProviderBuilder;
+import io.opentelemetry.sdk.trace.SpanProcessor;
+
+public class SamplingSpanProcessor implements SpanProcessor {
+
+  @Override
+  public void onStart(Context context, ReadWriteSpan readWriteSpan) {}
+
+  @Override
+  public boolean isStartRequired() {
+    return false;
+  }
+
+  @Override
+  public void onEnd(ReadableSpan readableSpan) {
+    DynamicSampler.evaluateSampled(readableSpan);
+  }
+
+  @Override
+  public boolean isEndRequired() {
+    return true;
+  }
+
+  public static SdkTracerProviderBuilder configure(
+      SdkTracerProviderBuilder sdkTracerProviderBuilder, ConfigProperties configProperties) {
+    new DynamicSampler(configProperties);
+    return sdkTracerProviderBuilder.addSpanProcessor(new SamplingSpanProcessor());
+  }
+}
diff --git a/custom/src/main/java/com/grafana/extensions/servertiming/ServerTimingHeaderCustomizer.java b/custom/src/main/java/com/grafana/extensions/servertiming/ServerTimingHeaderCustomizer.java
@@ -5,15 +5,15 @@
 
 package com.grafana.extensions.servertiming;
 
+import com.grafana.extensions.sampler.DynamicSampler;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.TraceFlags;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizer;
 import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
-import java.util.Set;
-import java.util.concurrent.ConcurrentSkipListSet;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
 
 /**
  * Adds {@code Server-Timing} header (and {@code Access-Control-Expose-Headers}) to the HTTP
@@ -26,8 +26,6 @@ public class ServerTimingHeaderCustomizer implements HttpServerResponseCustomize
   // not using volatile because this field is set only once during agent initialization
   static boolean enabled = false;
 
-  public static Set<String> sampledTraces = new ConcurrentSkipListSet<>();
-
   @Override
   public <RESPONSE> void customize(
       Context context, RESPONSE response, HttpServerResponseMutator<RESPONSE> responseMutator) {
@@ -40,18 +38,19 @@ public <RESPONSE> void customize(
   }
 
   static String toHeaderValue(Context context) {
-    SpanContext c = Span.fromContext(context).getSpanContext();
-    boolean sampled = sampledTraces.remove(c.getTraceId());
+    ReadWriteSpan span = (ReadWriteSpan) Span.fromContext(context);
+    SpanContext spanContext = span.getSpanContext();
+    boolean sampled = DynamicSampler.evaluateSampled(span);
     TraceParentHolder traceParentHolder = new TraceParentHolder();
     W3CTraceContextPropagator.getInstance()
         .inject(
             context.with(
                 Span.wrap(
                     SpanContext.create(
-                        c.getTraceId(),
-                        c.getSpanId(),
+                        spanContext.getTraceId(),
+                        spanContext.getSpanId(),
                         sampled ? TraceFlags.getSampled() : TraceFlags.getDefault(),
-                        c.getTraceState()))),
+                        spanContext.getTraceState()))),
             traceParentHolder,
             TraceParentHolder::set);
     return "traceparent;desc=\"" + traceParentHolder.traceParent + "\"";

diff --git a/custom/src/main/java/com/grafana/extensions/servertiming/ServerTimingHeaderReader.java b/custom/src/main/java/com/grafana/extensions/servertiming/ServerTimingHeaderReader.java
@@ -5,6 +5,7 @@
 
 package com.grafana.extensions.servertiming;
 
+import com.grafana.extensions.sampler.DynamicSampler;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
@@ -50,7 +51,7 @@ public <REQUEST, RESPONSE> void consume(
                         GETTER);
             SpanContext spanContext = Span.fromContext(traceparent).getSpanContext();
             if (spanContext.getTraceFlags().isSampled()) {
-              ServerTimingHeaderCustomizer.sampledTraces.add(spanContext.getTraceId());
+              DynamicSampler.setSampled(spanContext.getTraceId());
             }
           }
         }