remove dependency on abstract message in schema extractor (#7260)

DataDog · Oct 16, 2024 · a80d13e · a80d13e
1 parent 3f89f03
commit a80d13e
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 10 deletions.
diff --git a/...main/java/datadog/trace/instrumentation/protobuf_java/AbstractMessageInstrumentation.java b/...main/java/datadog/trace/instrumentation/protobuf_java/AbstractMessageInstrumentation.java
@@ -63,7 +63,11 @@ public static void onEnter(@Advice.This AbstractMessage message) {
       if (callDepth > 0) {
         return;
       }
-      SchemaExtractor.attachSchemaOnSpan(message, activeSpan(), SchemaExtractor.serialization);
+      if (message == null) {
+        return;
+      }
+      SchemaExtractor.attachSchemaOnSpan(
+          message.getDescriptorForType(), activeSpan(), SchemaExtractor.serialization);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

diff --git a/.../main/java/datadog/trace/instrumentation/protobuf_java/AbstractParserInstrumentation.java b/.../main/java/datadog/trace/instrumentation/protobuf_java/AbstractParserInstrumentation.java
@@ -80,7 +80,9 @@ public static void stopSpan(
       }
       if (message instanceof AbstractMessage) {
         SchemaExtractor.attachSchemaOnSpan(
-            (AbstractMessage) message, span, SchemaExtractor.deserialization);
+            ((AbstractMessage) message).getDescriptorForType(),
+            span,
+            SchemaExtractor.deserialization);
       }
     }
   }

diff --git a/...n/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/SchemaExtractor.java b/...n/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/SchemaExtractor.java
@@ -1,6 +1,5 @@
 package datadog.trace.instrumentation.protobuf_java;
 
-import com.google.protobuf.AbstractMessage;
 import com.google.protobuf.Descriptors;
 import com.google.protobuf.Descriptors.Descriptor;
 import com.google.protobuf.Descriptors.FieldDescriptor;
@@ -170,13 +169,6 @@ public void iterateOverSchema(SchemaBuilder builder) {
     extractSchema(descriptor, builder, 0);
   }
 
-  public static void attachSchemaOnSpan(AbstractMessage message, AgentSpan span, String operation) {
-    if (message == null) {
-      return;
-    }
-    attachSchemaOnSpan(message.getDescriptorForType(), span, operation);
-  }
-
   public static void attachSchemaOnSpan(
       Descriptors.Descriptor descriptor, AgentSpan span, String operation) {
     if (descriptor == null || span == null) {