Fix schema tracking for nested messages (#7690)

* fix schema tracking bug with recursive messages

* add generate script

* fmt

* add task for generating proto files

* remove old gradle

Author: piochelepiotr

dd-java-agent/instrumentation/protobuf/build.gradle

@@ -1,3 +1,9 @@
+plugins {
+  id "com.google.protobuf" version "0.9.4"
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
 muzzle {
   pass {
     group = "com.google.protobuf"
@@ -7,9 +13,30 @@ muzzle {
   }
 }
 
-apply from: "$rootDir/gradle/java.gradle"
+protobuf {
+  protoc {
+    def os = System.getProperty("os.name").toLowerCase()
+    def arch = System.getProperty("os.arch").toLowerCase()
+
+    // There is no m1 support for protoc 3.0.0, so require Rosetta
+    if (os.contains("mac") && arch.contains("aarch64")) {
+      artifact = "com.google.protobuf:protoc:3.0.0:osx-x86_64"
+    } else {
+      artifact = "com.google.protobuf:protoc:3.0.0"
+    }
+  }
+}
+
 
 dependencies {
   compileOnly group: 'com.google.protobuf', name: 'protobuf-java', version: '3.0.0'
   testImplementation group: 'com.google.protobuf', name: 'protobuf-java', version: '3.0.0'
 }
+
+sourceSets {
+  test {
+    java {
+      srcDir "$buildDir/generated/source/proto/test/java"
+    }
+  }
+}

dd-java-agent/instrumentation/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/AbstractMessageInstrumentation.java

@@ -10,6 +10,7 @@
 import com.google.protobuf.AbstractMessage;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.concurrent.ExecutionException;
 import net.bytebuddy.asm.Advice;
@@ -58,11 +59,19 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class WriteToAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnter(@Advice.This AbstractMessage message) {
+      final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(AbstractMessage.class);
+      if (callDepth > 0) {
+        return;
+      }
       SchemaExtractor.attachSchemaOnSpan(message, activeSpan(), SchemaExtractor.serialization);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(@Advice.Thrown final Throwable throwable) {
+      final int callDepth = CallDepthThreadLocalMap.decrementCallDepth(AbstractMessage.class);
+      if (callDepth > 0) {
+        return;
+      }
       AgentSpan span = activeSpan();
       if (throwable != null) {
         span.addThrowable(

dd-java-agent/instrumentation/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/AbstractParserInstrumentation.java

@@ -9,9 +9,11 @@
 
 import com.google.auto.service.AutoService;
 import com.google.protobuf.AbstractMessage;
+import com.google.protobuf.AbstractParser;
 import com.google.protobuf.MessageLite;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.concurrent.ExecutionException;
 import net.bytebuddy.asm.Advice;
@@ -55,9 +57,19 @@ public void methodAdvice(MethodTransformer transformer) {
   }
 
   public static class ParseFromAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void trackDepth() {
+      CallDepthThreadLocalMap.incrementCallDepth(AbstractParser.class);
+    }
+
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
         @Advice.Thrown final Throwable throwable, @Advice.Return MessageLite message) {
+      final int callDepth = CallDepthThreadLocalMap.decrementCallDepth(AbstractParser.class);
+      if (callDepth > 0) {
+        return;
+      }
       AgentSpan span = activeSpan();
       if (span == null) {
         return;

dd-java-agent/instrumentation/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/DynamicMessageInstrumentation.java

@@ -9,8 +9,10 @@
 
 import com.google.auto.service.AutoService;
 import com.google.protobuf.Descriptors.Descriptor;
+import com.google.protobuf.DynamicMessage;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import net.bytebuddy.asm.Advice;
 
 @AutoService(InstrumenterModule.class)
@@ -50,7 +52,16 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class ParseFromAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnter(@Advice.Argument(0) final Descriptor descriptor) {
+      final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(DynamicMessage.class);
+      if (callDepth > 0) {
+        return;
+      }
       SchemaExtractor.attachSchemaOnSpan(descriptor, activeSpan(), SchemaExtractor.deserialization);
     }
+
+    @Advice.OnMethodExit()
+    public static void trackDepth() {
+      CallDepthThreadLocalMap.decrementCallDepth(DynamicMessage.class);
+    }
   }
 }