remove dependency on abstract message in schema extractor

DataDog · Jun 28, 2024 · fec8d9f · fec8d9f
1 parent 08f0979
commit fec8d9f
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
@@ -58,7 +58,11 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class WriteToAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnter(@Advice.This AbstractMessage message) {
-      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
@@ -68,7 +68,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) {