Capture Kafka consumer group membership on join

dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/ConsumerCoordinatorInstrumentation.java

@@ -10,6 +10,7 @@
 import datadog.trace.api.datastreams.DataStreamsTags;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.instrumentation.kafka_common.KafkaConfigHelper;
 import datadog.trace.instrumentation.kafka_common.MetadataState;
 import java.util.HashMap;
 import java.util.Map;
@@ -61,6 +62,7 @@ public String instrumentedType() {
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".KafkaConsumerInfo",
+      "datadog.trace.instrumentation.kafka_common.KafkaConfigHelper",
       "datadog.trace.instrumentation.kafka_common.PendingConfig",
       "datadog.trace.instrumentation.kafka_common.MetadataState",
     };
@@ -71,6 +73,9 @@ public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         isMethod().and(named("sendOffsetCommitRequest")).and(takesArguments(1)),
         ConsumerCoordinatorInstrumentation.class.getName() + "$CommitOffsetAdvice");
+    transformer.applyAdvice(
+        isMethod().and(named("onJoinComplete")).and(takesArguments(4)),
+        ConsumerCoordinatorInstrumentation.class.getName() + "$JoinGroupAdvice");
   }
 
   public static class CommitOffsetAdvice {
@@ -127,4 +132,44 @@ public static void muzzleCheck(ConsumerRecord record) {
       record.headers();
     }
   }
+
+  public static class JoinGroupAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void trackJoinGroup(
+        @Advice.This ConsumerCoordinator coordinator,
+        @Advice.Argument(0) final int generationId,
+        @Advice.Argument(1) final String memberId,
+        @Advice.Argument(2) final String memberProtocol) {
+      if (memberId == null || memberId.isEmpty()) {
+        return;
+      }
+      KafkaConsumerInfo kafkaConsumerInfo =
+          InstrumentationContext.get(ConsumerCoordinator.class, KafkaConsumerInfo.class)
+              .get(coordinator);
+      if (kafkaConsumerInfo == null) {
+        return;
+      }
+      if (memberId.equals(kafkaConsumerInfo.getLastReportedMemberId())
+          && generationId == kafkaConsumerInfo.getLastReportedGenerationId()) {
+        return;
+      }
+
+      String consumerGroup = kafkaConsumerInfo.getConsumerGroup();
+      Metadata consumerMetadata = kafkaConsumerInfo.getClientMetadata();
+      String clusterId = null;
+      if (consumerMetadata != null) {
+        MetadataState metadataState =
+            InstrumentationContext.get(Metadata.class, MetadataState.class).get(consumerMetadata);
+        clusterId = metadataState != null ? metadataState.clusterId : null;
+      }
+      if (KafkaConfigHelper.reportConsumerGroupMember(
+          clusterId, consumerGroup, memberId, generationId, memberProtocol)) {
+        kafkaConsumerInfo.setLastReportedMembership(memberId, generationId);
+      }
+    }
+
+    public static void muzzleCheck(ConsumerRecord record) {
+      record.headers();
+    }
+  }
 }

dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfo.java

@@ -8,6 +8,8 @@ public class KafkaConsumerInfo {
   private final String consumerGroup;
   private final Metadata clientMetadata;
   private final String bootstrapServers;
+  private volatile String lastReportedMemberId;
+  private volatile int lastReportedGenerationId = Integer.MIN_VALUE;
 
   public KafkaConsumerInfo(String consumerGroup, Metadata clientMetadata, String bootstrapServers) {
     this.consumerGroup = consumerGroup;
@@ -36,6 +38,20 @@ public String getBootstrapServers() {
     return bootstrapServers;
   }
 
+  @Nullable
+  public String getLastReportedMemberId() {
+    return lastReportedMemberId;
+  }
+
+  public int getLastReportedGenerationId() {
+    return lastReportedGenerationId;
+  }
+
+  public void setLastReportedMembership(String memberId, int generationId) {
+    this.lastReportedMemberId = memberId;
+    this.lastReportedGenerationId = generationId;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) return true;

dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java/datadog/trace/instrumentation/kafka_clients38/ConsumerCoordinatorInstrumentation.java

@@ -45,6 +45,7 @@ public String instrumentedType() {
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".KafkaConsumerInfo",
+      "datadog.trace.instrumentation.kafka_common.KafkaConfigHelper",
       "datadog.trace.instrumentation.kafka_common.PendingConfig",
       "datadog.trace.instrumentation.kafka_common.MetadataState",
     };
@@ -55,5 +56,8 @@ public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         isMethod().and(named("sendOffsetCommitRequest")).and(takesArguments(1)),
         packageName + ".ConsumerCoordinatorAdvice");
+    transformer.applyAdvice(
+        isMethod().and(named("onJoinComplete")).and(takesArguments(4)),
+        packageName + ".JoinGroupAdvice");
   }
 }

dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/JoinGroupAdvice.java

@@ -0,0 +1,49 @@
+package datadog.trace.instrumentation.kafka_clients38;
+
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.instrumentation.kafka_common.KafkaConfigHelper;
+import datadog.trace.instrumentation.kafka_common.MetadataState;
+import net.bytebuddy.asm.Advice;
+import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.internals.ConsumerCoordinator;
+
+public class JoinGroupAdvice {
+  @Advice.OnMethodExit(suppress = Throwable.class)
+  public static void trackJoinGroup(
+      @Advice.This ConsumerCoordinator coordinator,
+      @Advice.Argument(0) final int generationId,
+      @Advice.Argument(1) final String memberId,
+      @Advice.Argument(2) final String memberProtocol) {
+    if (memberId == null || memberId.isEmpty()) {
+      return;
+    }
+    KafkaConsumerInfo kafkaConsumerInfo =
+        InstrumentationContext.get(ConsumerCoordinator.class, KafkaConsumerInfo.class)
+            .get(coordinator);
+    if (kafkaConsumerInfo == null) {
+      return;
+    }
+    if (memberId.equals(kafkaConsumerInfo.getLastReportedMemberId().orElse(null))
+        && generationId == kafkaConsumerInfo.getLastReportedGenerationId()) {
+      return;
+    }
+
+    String consumerGroup = kafkaConsumerInfo.getConsumerGroup().orElse(null);
+    Metadata consumerMetadata = kafkaConsumerInfo.getmetadata().orElse(null);
+    String clusterId = null;
+    if (consumerMetadata != null) {
+      MetadataState metadataState =
+          InstrumentationContext.get(Metadata.class, MetadataState.class).get(consumerMetadata);
+      clusterId = metadataState != null ? metadataState.clusterId : null;
+    }
+    if (KafkaConfigHelper.reportConsumerGroupMember(
+        clusterId, consumerGroup, memberId, generationId, memberProtocol)) {
+      kafkaConsumerInfo.setLastReportedMembership(memberId, generationId);
+    }
+  }
+
+  public static void muzzleCheck(ConsumerRecord record) {
+    record.headers();
+  }
+}

dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/KafkaConsumerInfo.java

@@ -8,6 +8,8 @@ public class KafkaConsumerInfo {
   private final String consumerGroup;
   private final Metadata metadata;
   private final String bootstrapServers;
+  private volatile String lastReportedMemberId;
+  private volatile int lastReportedGenerationId = Integer.MIN_VALUE;
 
   public KafkaConsumerInfo(String consumerGroup, Metadata metadata, String bootstrapServers) {
     this.consumerGroup = consumerGroup;
@@ -33,6 +35,19 @@ public Optional<String> getBootstrapServers() {
     return Optional.ofNullable(bootstrapServers);
   }
 
+  public Optional<String> getLastReportedMemberId() {
+    return Optional.ofNullable(lastReportedMemberId);
+  }
+
+  public int getLastReportedGenerationId() {
+    return lastReportedGenerationId;
+  }
+
+  public void setLastReportedMembership(String memberId, int generationId) {
+    this.lastReportedMemberId = memberId;
+    this.lastReportedGenerationId = generationId;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) return true;

dd-java-agent/instrumentation/kafka/kafka-common/src/main/java/datadog/trace/instrumentation/kafka_common/KafkaConfigHelper.java

@@ -105,6 +105,28 @@ public static void storePendingConsumerConfig(
     log.debug("Stored pending consumer config (cluster ID not yet known)");
   }
 
+  public static boolean reportConsumerGroupMember(
+      String clusterId,
+      String consumerGroup,
+      String memberId,
+      int generationId,
+      String memberProtocol) {
+    if (memberId == null || memberId.isEmpty() || clusterId == null || clusterId.isEmpty()) {
+      return false;
+    }
+    if (Config.get().isDataStreamsEnabled()) {
+      AgentTracer.get()
+          .getDataStreamsMonitoring()
+          .reportKafkaConsumerGroupMember(
+              clusterId,
+              consumerGroup != null ? consumerGroup : "",
+              memberId,
+              generationId,
+              memberProtocol != null ? memberProtocol : "");
+    }
+    return true;
+  }
+
   /** Called from metadata update advice when the cluster ID becomes available. */
   public static void reportPendingConfig(MetadataState state, String clusterId) {
     PendingConfig pending = state.takePendingConfig();

dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java

@@ -305,6 +305,26 @@ public void reportKafkaConfig(
             getThreadServiceName()));
   }
 
+  @Override
+  public void reportKafkaConsumerGroupMember(
+      String kafkaClusterId,
+      String consumerGroup,
+      String memberId,
+      int generationId,
+      String memberProtocol) {
+    inbox.offer(
+        new KafkaConfigReport(
+            "kafka_consumer",
+            kafkaClusterId,
+            consumerGroup,
+            memberId,
+            generationId,
+            memberProtocol,
+            Collections.<String, String>emptyMap(),
+            timeSource.getCurrentTimeNanos(),
+            getThreadServiceName()));
+  }
+
   @Override
   public void setCheckpoint(AgentSpan span, DataStreamsContext context) {
     PathwayContext pathwayContext = span.spanContext().getPathwayContext();

dd-trace-core/src/main/java/datadog/trace/core/datastreams/MsgPackDatastreamsPayloadWriter.java

@@ -57,6 +57,9 @@ public class MsgPackDatastreamsPayloadWriter implements DatastreamsPayloadWriter
   private static final byte[] CONFIG_TYPE = "Type".getBytes(ISO_8859_1);
   private static final byte[] CONFIG_KAFKA_CLUSTER_ID = "KafkaClusterId".getBytes(ISO_8859_1);
   private static final byte[] CONFIG_CONSUMER_GROUP = "ConsumerGroup".getBytes(ISO_8859_1);
+  private static final byte[] CONFIG_MEMBER_ID = "MemberId".getBytes(ISO_8859_1);
+  private static final byte[] CONFIG_GENERATION_ID = "GenerationId".getBytes(ISO_8859_1);
+  private static final byte[] CONFIG_MEMBER_PROTOCOL = "MemberProtocol".getBytes(ISO_8859_1);
   private static final byte[] CONFIG_ENTRIES = "Config".getBytes(ISO_8859_1);
 
   private static final int INITIAL_CAPACITY = 512 * 1024;
@@ -290,7 +293,8 @@ private void writeKafkaConfigs(List<KafkaConfigReport> configs, Writable packer)
     packer.writeUTF8(CONFIGS);
     packer.startArray(configs.size());
     for (KafkaConfigReport config : configs) {
-      packer.startMap(4); // Type, KafkaClusterId, ConsumerGroup, Config
+      // Type, KafkaClusterId, ConsumerGroup, MemberId, GenerationId, MemberProtocol, Config
+      packer.startMap(7);
 
       packer.writeUTF8(CONFIG_TYPE);
       packer.writeString(config.getType(), null);
@@ -301,6 +305,15 @@ private void writeKafkaConfigs(List<KafkaConfigReport> configs, Writable packer)
       packer.writeUTF8(CONFIG_CONSUMER_GROUP);
       packer.writeString(config.getConsumerGroup(), null);
 
+      packer.writeUTF8(CONFIG_MEMBER_ID);
+      packer.writeString(config.getMemberId(), null);
+
+      packer.writeUTF8(CONFIG_GENERATION_ID);
+      packer.writeInt(config.getGenerationId());
+
+      packer.writeUTF8(CONFIG_MEMBER_PROTOCOL);
+      packer.writeString(config.getMemberProtocol(), null);
+
       packer.writeUTF8(CONFIG_ENTRIES);
       Map<String, String> entries = config.getConfig();
       packer.startMap(entries.size());