Suppress instrumentation based on suppress Context key (#9739)

open-telemetry · Nov 17, 2023 · 683c311 · 683c311
1 parent c5cb948
commit 683c311
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 1 deletion.
diff --git a/instrumentation-api/build.gradle.kts b/instrumentation-api/build.gradle.kts
@@ -20,6 +20,7 @@ dependencies {
 
   testImplementation(project(":testing-common"))
   testImplementation("io.opentelemetry:opentelemetry-sdk-testing")
+  testImplementation("io.opentelemetry:opentelemetry-exporter-common")
   testImplementation("org.junit-pioneer:junit-pioneer")
 
   jmhImplementation(project(":instrumentation-api-semconv"))

diff --git a/.../src/main/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterBuilder.java b/.../src/main/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterBuilder.java
@@ -354,7 +354,8 @@ private String getSchemaUrl() {
   }
 
   SpanSuppressor buildSpanSuppressor() {
-    return spanSuppressionStrategy.create(getSpanKeysFromAttributesExtractors());
+    return new SpanSuppressors.ByContextKey(
+        spanSuppressionStrategy.create(getSpanKeysFromAttributesExtractors()));
   }
 
   private Set<SpanKey> getSpanKeysFromAttributesExtractors() {

diff --git a/...-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/SpanSuppressors.java b/...-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/SpanSuppressors.java
@@ -10,6 +10,8 @@
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.internal.SpanKey;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.util.Map;
 import java.util.Set;
 
@@ -85,4 +87,48 @@ public boolean shouldSuppress(Context parentContext, SpanKind spanKind) {
       return true;
     }
   }
+
+  static class ByContextKey implements SpanSuppressor {
+    private final SpanSuppressor delegate;
+    private final Method shouldSuppressInstrumentation;
+
+    ByContextKey(SpanSuppressor delegate) {
+      this.delegate = delegate;
+      Method shouldSuppressInstrumentation;
+      try {
+        Class<?> instrumentationUtil =
+            Class.forName("io.opentelemetry.exporter.internal.InstrumentationUtil");
+        shouldSuppressInstrumentation =
+            instrumentationUtil.getDeclaredMethod("shouldSuppressInstrumentation", Context.class);
+      } catch (ClassNotFoundException | NoSuchMethodException e) {
+        shouldSuppressInstrumentation = null;
+      }
+      this.shouldSuppressInstrumentation = shouldSuppressInstrumentation;
+    }
+
+    @Override
+    public Context storeInContext(Context context, SpanKind spanKind, Span span) {
+      return delegate.storeInContext(context, spanKind, span);
+    }
+
+    @Override
+    public boolean shouldSuppress(Context parentContext, SpanKind spanKind) {
+      if (suppressByContextKey(parentContext)) {
+        return true;
+      }
+      return delegate.shouldSuppress(parentContext, spanKind);
+    }
+
+    private boolean suppressByContextKey(Context context) {
+      if (shouldSuppressInstrumentation == null) {
+        return false;
+      }
+
+      try {
+        return (boolean) shouldSuppressInstrumentation.invoke(null, context);
+      } catch (IllegalAccessException | InvocationTargetException e) {
+        return false;
+      }
+    }
+  }
 }
diff --git a/...t/java/io/opentelemetry/instrumentation/api/instrumenter/SpanSuppressionStrategyTest.java b/...t/java/io/opentelemetry/instrumentation/api/instrumenter/SpanSuppressionStrategyTest.java
@@ -16,6 +16,7 @@
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.exporter.internal.InstrumentationUtil;
 import io.opentelemetry.instrumentation.api.internal.SpanKey;
 import java.util.HashSet;
 import java.util.Set;
@@ -157,4 +158,24 @@ void semconv_shouldNotSuppressContextWithPartiallyDifferentSpanKeys() {
 
     assertFalse(suppressor.shouldSuppress(context, SpanKind.SERVER));
   }
+
+  @Test
+  void context_shouldSuppressWhenKeyIsAvailableAndTrue() {
+    InstrumentationUtil.suppressInstrumentation(
+        () -> {
+          SpanSuppressor suppressor =
+              new SpanSuppressors.ByContextKey(SpanSuppressionStrategy.NONE.create(emptySet()));
+
+          assertTrue(suppressor.shouldSuppress(Context.current(), SpanKind.CLIENT));
+        });
+  }
+
+  @Test
+  void context_shouldNotSuppressWhenKeyIsNotAvailable() {
+    Context context = Context.current();
+    SpanSuppressor suppressor =
+        new SpanSuppressors.ByContextKey(SpanSuppressionStrategy.NONE.create(emptySet()));
+
+    assertFalse(suppressor.shouldSuppress(context, SpanKind.CLIENT));
+  }
 }