updated following review

DataDog · Jun 27, 2024 · dc5f4fb · dc5f4fb
1 parent c8be173
commit dc5f4fb
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 20 deletions.
diff --git a/dd-java-agent/instrumentation/avro/build.gradle b/dd-java-agent/instrumentation/avro/build.gradle
@@ -9,13 +9,18 @@ muzzle {
 
 apply from: "$rootDir/gradle/java.gradle"
 
+addTestSuiteForDir('latestDepTest','test')
+
 dependencies {
   compileOnly group: 'org.apache.avro', name: 'avro', version: '1.11.3'
   testImplementation group: 'org.apache.avro', name: 'avro', version: '1.11.3'
+  latestDepTestImplementation group: 'org.apache.avro', name: 'avro', version: '1.+'
 
   compileOnly group: 'org.apache.hadoop', name: 'hadoop-common', version: '3.3.1'
   testImplementation group: 'org.apache.hadoop', name: 'hadoop-common', version: '3.3.1'
+  latestDepTestImplementation group: 'org.apache.hadoop', name: 'hadoop-common', version: '3.+'
 
   compileOnly group: 'org.apache.avro', name: 'avro-mapred', version: '1.11.3'
   testImplementation group: 'org.apache.avro', name: 'avro-mapred', version: '1.11.3'
+  latestDepTestImplementation group: 'org.apache.avro', name: 'avro-mapred', version: '1.11.3'
 }
diff --git a/...o/src/main/java/datadog/trace/instrumentation/avro/GenericDatumReaderInstrumentation.java b/...o/src/main/java/datadog/trace/instrumentation/avro/GenericDatumReaderInstrumentation.java
@@ -1,13 +1,14 @@
 package datadog.trace.instrumentation.avro;
 
-import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -18,17 +19,14 @@
 public final class GenericDatumReaderInstrumentation extends InstrumenterModule.Tracing
     implements Instrumenter.ForTypeHierarchy {
 
-  static final String instrumentationName = "avro";
-  static final String TARGET_TYPE = "org.apache.avro.generic.GenericDatumReader";
-
   public GenericDatumReaderInstrumentation() {
-    super(instrumentationName);
+    super("avro");
   }
 
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".SchemaExtractor", "datadog.trace.instrumentation.avro.SchemaExtractor$1",
+      packageName + ".SchemaExtractor", packageName + ".SchemaExtractor$1",
     };
   }
 
@@ -41,24 +39,27 @@ public void methodAdvice(MethodTransformer transformer) {
 
   @Override
   public String hierarchyMarkerType() {
-    return TARGET_TYPE;
+    return "org.apache.avro.generic.GenericDatumReader";
   }
 
   @Override
   public ElementMatcher<TypeDescription> hierarchyMatcher() {
-    return hasSuperType(named(hierarchyMarkerType()));
+    return extendsClass(named(hierarchyMarkerType()));
   }
 
   public static class GenericDatumReaderAdvice {
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(@Advice.This GenericDatumReader reader) {
+      if (!Config.get().isDataStreamsEnabled()) {
+        return;
+      }
       AgentSpan span = activeSpan();
       if (span == null) {
         return;
       }
       if (reader != null) {
         SchemaExtractor.attachSchemaOnSpan(
-            reader.getSchema(), span, SchemaExtractor.deserialization);
+            reader.getSchema(), span, SchemaExtractor.DESERIALIZATION);
       }
     }
   }

diff --git a/...o/src/main/java/datadog/trace/instrumentation/avro/GenericDatumWriterInstrumentation.java b/...o/src/main/java/datadog/trace/instrumentation/avro/GenericDatumWriterInstrumentation.java
@@ -1,6 +1,6 @@
 package datadog.trace.instrumentation.avro;
 
-import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -9,6 +9,7 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -19,17 +20,14 @@
 public final class GenericDatumWriterInstrumentation extends InstrumenterModule.Tracing
     implements Instrumenter.ForTypeHierarchy {
 
-  static final String instrumentationName = "avro";
-  static final String TARGET_TYPE = "org.apache.avro.generic.GenericDatumWriter";
-
   public GenericDatumWriterInstrumentation() {
-    super(instrumentationName);
+    super("avro");
   }
 
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".SchemaExtractor", "datadog.trace.instrumentation.avro.SchemaExtractor$1",
+      packageName + ".SchemaExtractor", packageName + ".SchemaExtractor$1",
     };
   }
 
@@ -42,24 +40,27 @@ public void methodAdvice(MethodTransformer transformer) {
 
   @Override
   public String hierarchyMarkerType() {
-    return TARGET_TYPE;
+    return "org.apache.avro.generic.GenericDatumWriter";
   }
 
   @Override
   public ElementMatcher<TypeDescription> hierarchyMatcher() {
-    return hasSuperType(named(hierarchyMarkerType()));
+    return extendsClass(named(hierarchyMarkerType()));
   }
 
   public static class GenericDatumWriterAdvice {
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(@Advice.FieldValue("root") Schema root) {
+      if (!Config.get().isDataStreamsEnabled()) {
+        return;
+      }
       AgentSpan span = activeSpan();
       if (span == null) {
         return;
       }
 
       if (root != null) {
-        SchemaExtractor.attachSchemaOnSpan(root, span, SchemaExtractor.serialization);
+        SchemaExtractor.attachSchemaOnSpan(root, span, SchemaExtractor.SERIALIZATION);
       }
     }
   }

diff --git a/...nstrumentation/avro/src/main/java/datadog/trace/instrumentation/avro/SchemaExtractor.java b/...nstrumentation/avro/src/main/java/datadog/trace/instrumentation/avro/SchemaExtractor.java
@@ -11,8 +11,8 @@
 import org.apache.avro.Schema.Field;
 
 public class SchemaExtractor implements SchemaIterator {
-  public static final String serialization = "serialization";
-  public static final String deserialization = "deserialization";
+  public static final String SERIALIZATION = "serialization";
+  public static final String DESERIALIZATION = "deserialization";
   private static final String AVRO = "avro";
 
   private final org.apache.avro.Schema schema;