Track schema registry usage

.gitignore

@@ -46,6 +46,10 @@ out/
 ######################
 .vscode
 
+# Cursor #
+##########
+.cursor
+
 # Others #
 ##########
 /logs/*

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/datastreams/RecordingDatastreamsPayloadWriter.groovy

@@ -19,6 +19,9 @@ class RecordingDatastreamsPayloadWriter implements DatastreamsPayloadWriter {
   @SuppressWarnings('UnusedPrivateField')
   private final Set<DataStreamsTags> backlogs = []
 
+  @SuppressWarnings('UnusedPrivateField')
+  private final List<StatsBucket.SchemaKey> schemaRegistryUsages = []
+
   private final Set<String> serviceNameOverrides = []
 
   @Override
@@ -33,6 +36,11 @@ class RecordingDatastreamsPayloadWriter implements DatastreamsPayloadWriter {
           this.@backlogs.add(backlog.getKey())
         }
       }
+      if (bucket.schemaRegistryUsages != null) {
+        for (Map.Entry<StatsBucket.SchemaKey, Long> usage : bucket.schemaRegistryUsages) {
+          this.@schemaRegistryUsages.add(usage.getKey())
+        }
+      }
     }
   }
 
@@ -52,10 +60,15 @@ class RecordingDatastreamsPayloadWriter implements DatastreamsPayloadWriter {
     Collections.unmodifiableList(new ArrayList<>(this.@backlogs))
   }
 
+  synchronized List<StatsBucket.SchemaKey> getSchemaRegistryUsages() {
+    Collections.unmodifiableList(new ArrayList<>(this.@schemaRegistryUsages))
+  }
+
   synchronized void clear() {
     this.@payloads.clear()
     this.@groups.clear()
     this.@backlogs.clear()
+    this.@schemaRegistryUsages.clear()
   }
 
   void waitForPayloads(int count, long timeout = TimeUnit.SECONDS.toMillis(3)) {
@@ -70,6 +83,10 @@ class RecordingDatastreamsPayloadWriter implements DatastreamsPayloadWriter {
     waitFor(count, timeout, this.@backlogs)
   }
 
+  void waitForSchemaRegistryUsages(int count, long timeout = TimeUnit.SECONDS.toMillis(3)) {
+    waitFor(count, timeout, this.@schemaRegistryUsages)
+  }
+
   private static void waitFor(int count, long timeout, Collection collection) {
     long deadline = System.currentTimeMillis() + timeout
     while (System.currentTimeMillis() < deadline) {

dd-java-agent/instrumentation/confluent-schema-registry/confluent-schema-registry-7.0/build.gradle

@@ -0,0 +1,26 @@
+apply from: "$rootDir/gradle/java.gradle"
+
+muzzle {
+  pass {
+    group = "io.confluent"
+    module = "kafka-schema-registry-client"
+    versions = "[7.0.0,)"
+    assertInverse = true
+  }
+}
+
+dependencies {
+  compileOnly project(':dd-java-agent:instrumentation:kafka:kafka-common')
+  compileOnly group: 'io.confluent', name: 'kafka-schema-registry-client', version: '7.0.0'
+  compileOnly group: 'io.confluent', name: 'kafka-avro-serializer', version: '7.0.0'
+  compileOnly group: 'io.confluent', name: 'kafka-protobuf-serializer', version: '7.0.0'
+  compileOnly group: 'org.apache.kafka', name: 'kafka-clients', version: '3.0.0'
+
+  testImplementation project(':dd-java-agent:instrumentation:kafka:kafka-common')
+  testImplementation group: 'io.confluent', name: 'kafka-schema-registry-client', version: '7.5.2'
+  testImplementation group: 'io.confluent', name: 'kafka-avro-serializer', version: '7.5.2'
+  testImplementation group: 'io.confluent', name: 'kafka-protobuf-serializer', version: '7.5.1'
+  testImplementation group: 'org.apache.kafka', name: 'kafka-clients', version: '3.5.0'
+  testImplementation group: 'org.apache.avro', name: 'avro', version: '1.11.0'
+}
+

dd-java-agent/instrumentation/confluent-schema-registry/confluent-schema-registry-7.0/src/main/java/datadog/trace/instrumentation/confluentschemaregistry/KafkaDeserializerInstrumentation.java

@@ -0,0 +1,115 @@
+package datadog.trace.instrumentation.confluentschemaregistry;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.Instrumenter.MethodTransformer;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.instrumentation.kafka_common.ClusterIdHolder;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import org.apache.kafka.common.serialization.Deserializer;
+
+/**
+ * Instruments Confluent Schema Registry deserializers (Avro, Protobuf, and JSON) to capture
+ * deserialization operations.
+ */
+@AutoService(InstrumenterModule.class)
+public class KafkaDeserializerInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  public KafkaDeserializerInstrumentation() {
+    super("confluent-schema-registry", "kafka");
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "io.confluent.kafka.serializers.KafkaAvroDeserializer",
+      "io.confluent.kafka.serializers.json.KafkaJsonSchemaDeserializer",
+      "io.confluent.kafka.serializers.protobuf.KafkaProtobufDeserializer"
+    };
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.instrumentation.kafka_common.ClusterIdHolder",
+      packageName + ".SchemaIdExtractor"
+    };
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    Map<String, String> contextStores = new HashMap<>();
+    contextStores.put("org.apache.kafka.common.serialization.Deserializer", "java.lang.Boolean");
+    return contextStores;
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    // Instrument configure to capture isKey value
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("configure"))
+            .and(isPublic())
+            .and(takesArguments(2))
+            .and(takesArgument(1, boolean.class)),
+        getClass().getName() + "$ConfigureAdvice");
+
+    // Instrument deserialize(String topic, Headers headers, byte[] data)
+    // The 2-arg version calls this one, so we only need to instrument this to avoid duplicates
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("deserialize"))
+            .and(isPublic())
+            .and(takesArguments(3))
+            .and(takesArgument(0, String.class))
+            .and(takesArgument(2, byte[].class)),
+        getClass().getName() + "$DeserializeAdvice");
+  }
+
+  public static class ConfigureAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(
+        @Advice.This Deserializer deserializer, @Advice.Argument(1) boolean isKey) {
+      // Store the isKey value in InstrumentationContext for later use
+      InstrumentationContext.get(Deserializer.class, Boolean.class).put(deserializer, isKey);
+    }
+  }
+
+  public static class DeserializeAdvice {
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.This Deserializer deserializer,
+        @Advice.Argument(0) String topic,
+        @Advice.Argument(2) byte[] data,
+        @Advice.Return Object result,
+        @Advice.Thrown Throwable throwable) {
+
+      // Get isKey from InstrumentationContext (stored during configure)
+      Boolean isKeyObj =
+          InstrumentationContext.get(Deserializer.class, Boolean.class).get(deserializer);
+      boolean isKey = isKeyObj != null && isKeyObj;
+
+      // Get cluster ID from thread-local (set by Kafka consumer instrumentation)
+      String clusterId = ClusterIdHolder.get();
+
+      boolean isSuccess = throwable == null;
+      int schemaId = isSuccess ? SchemaIdExtractor.extractSchemaId(data) : -1;
+
+      // Record the schema registry usage
+      AgentTracer.get()
+          .getDataStreamsMonitoring()
+          .reportSchemaRegistryUsage(topic, clusterId, schemaId, isSuccess, isKey, "deserialize");
+    }
+  }
+}

dd-java-agent/instrumentation/confluent-schema-registry/confluent-schema-registry-7.0/src/main/java/datadog/trace/instrumentation/confluentschemaregistry/KafkaSerializerInstrumentation.java

@@ -0,0 +1,114 @@
+package datadog.trace.instrumentation.confluentschemaregistry;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.Instrumenter.MethodTransformer;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.instrumentation.kafka_common.ClusterIdHolder;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import org.apache.kafka.common.serialization.Serializer;
+
+/**
+ * Instruments Confluent Schema Registry serializers (Avro, Protobuf, and JSON) to capture
+ * serialization operations.
+ */
+@AutoService(InstrumenterModule.class)
+public class KafkaSerializerInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  public KafkaSerializerInstrumentation() {
+    super("confluent-schema-registry", "kafka");
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "io.confluent.kafka.serializers.KafkaAvroSerializer",
+      "io.confluent.kafka.serializers.json.KafkaJsonSchemaSerializer",
+      "io.confluent.kafka.serializers.protobuf.KafkaProtobufSerializer"
+    };
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.instrumentation.kafka_common.ClusterIdHolder",
+      packageName + ".SchemaIdExtractor"
+    };
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    Map<String, String> contextStores = new HashMap<>();
+    contextStores.put("org.apache.kafka.common.serialization.Serializer", "java.lang.Boolean");
+    return contextStores;
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    // Instrument configure to capture isKey value
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("configure"))
+            .and(isPublic())
+            .and(takesArguments(2))
+            .and(takesArgument(1, boolean.class)),
+        getClass().getName() + "$ConfigureAdvice");
+
+    // Instrument both serialize(String topic, Object data)
+    // and serialize(String topic, Headers headers, Object data) for Kafka 2.1+
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("serialize"))
+            .and(isPublic())
+            .and(takesArgument(0, String.class))
+            .and(returns(byte[].class)),
+        getClass().getName() + "$SerializeAdvice");
+  }
+
+  public static class ConfigureAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(
+        @Advice.This Serializer serializer, @Advice.Argument(1) boolean isKey) {
+      // Store the isKey value in InstrumentationContext for later use
+      InstrumentationContext.get(Serializer.class, Boolean.class).put(serializer, isKey);
+    }
+  }
+
+  public static class SerializeAdvice {
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.This Serializer serializer,
+        @Advice.Argument(0) String topic,
+        @Advice.Return byte[] result,
+        @Advice.Thrown Throwable throwable) {
+
+      // Get isKey from InstrumentationContext (stored during configure)
+      Boolean isKeyObj =
+          InstrumentationContext.get(Serializer.class, Boolean.class).get(serializer);
+      boolean isKey = isKeyObj != null && isKeyObj;
+
+      // Get cluster ID from thread-local (set by Kafka producer instrumentation)
+      String clusterId = ClusterIdHolder.get();
+
+      boolean isSuccess = throwable == null;
+      int schemaId = isSuccess ? SchemaIdExtractor.extractSchemaId(result) : -1;
+
+      // Record the schema registry usage
+      AgentTracer.get()
+          .getDataStreamsMonitoring()
+          .reportSchemaRegistryUsage(topic, clusterId, schemaId, isSuccess, isKey, "serialize");
+    }
+  }
+}

dd-java-agent/instrumentation/confluent-schema-registry/confluent-schema-registry-7.0/src/main/java/datadog/trace/instrumentation/confluentschemaregistry/SchemaIdExtractor.java

@@ -0,0 +1,23 @@
+package datadog.trace.instrumentation.confluentschemaregistry;
+
+/**
+ * Helper class to extract schema ID from Confluent Schema Registry wire format. Wire format:
+ * [magic_byte][4-byte schema id][data]
+ */
+public class SchemaIdExtractor {
+  public static int extractSchemaId(byte[] data) {
+    if (data == null || data.length < 5 || data[0] != 0) {
+      return -1;
+    }
+
+    try {
+      // Confluent wire format: [magic_byte][4-byte schema id][data]
+      return ((data[1] & 0xFF) << 24)
+          | ((data[2] & 0xFF) << 16)
+          | ((data[3] & 0xFF) << 8)
+          | (data[4] & 0xFF);
+    } catch (Throwable ignored) {
+      return -1;
+    }
+  }
+}