[WIP - DO NOT REVIEW] Add DD_TRACE_KAFKA_CREATE_CONSUMER_SCOPE_ENABLED for Kafka consumers

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

@@ -1,6 +1,8 @@
 package datadog.trace.instrumentation.kafka_clients;
 
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nullable;
 import org.apache.kafka.clients.Metadata;
 
@@ -9,6 +11,10 @@ public class KafkaConsumerInfo {
   private final Metadata clientMetadata;
   private final String bootstrapServers;
 
+  // handle to the consume span this consumer left lingering past its last poll loop; not part of
+  // consumer identity, so excluded from equals/hashCode
+  private final AtomicReference<AgentSpan> deferredConsumeSpan = new AtomicReference<>();
+
   public KafkaConsumerInfo(String consumerGroup, Metadata clientMetadata, String bootstrapServers) {
     this.consumerGroup = consumerGroup;
     this.clientMetadata = clientMetadata;
@@ -21,6 +27,14 @@ public KafkaConsumerInfo(String consumerGroup, String bootstrapServers) {
     this.bootstrapServers = bootstrapServers;
   }
 
+  public void setDeferredConsumeSpan(AgentSpan span) {
+    deferredConsumeSpan.set(span);
+  }
+
+  public AgentSpan getAndClearDeferredConsumeSpan() {
+    return deferredConsumeSpan.getAndSet(null);
+  }
+
   @Nullable
   public String getConsumerGroup() {
     return consumerGroup;

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

@@ -87,6 +87,7 @@ public String[] helperClassNames() {
     return new String[] {
       packageName + ".KafkaDecorator",
       packageName + ".KafkaConsumerInfo",
+      packageName + ".KafkaConsumerInstrumentationHelper",
       "datadog.trace.instrumentation.kafka_common.ClusterIdHolder",
       "datadog.trace.instrumentation.kafka_common.KafkaConfigHelper",
       "datadog.trace.instrumentation.kafka_common.PendingConfig",
@@ -120,6 +121,15 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(takesArguments(1))
             .and(returns(named("org.apache.kafka.clients.consumer.ConsumerRecords"))),
         KafkaConsumerInfoInstrumentation.class.getName() + "$RecordsAdvice");
+
+    // close()/unsubscribe() end the poll loop for good, so a consume span deferred past the last
+    // poll must be finished here rather than left dangling.
+    transformer.applyAdvice(
+        isMethod().and(isPublic()).and(named("close")),
+        KafkaConsumerInfoInstrumentation.class.getName() + "$CloseScopeAdvice");
+    transformer.applyAdvice(
+        isMethod().and(isPublic()).and(named("unsubscribe")).and(takesArguments(0)),
+        KafkaConsumerInfoInstrumentation.class.getName() + "$CloseScopeAdvice");
   }
 
   public static class ConstructorAdviceNot27 {
@@ -237,9 +247,14 @@ public static void muzzleCheck(ConsumerRecord record) {
   public static class RecordsAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope onEnter(@Advice.This KafkaConsumer consumer) {
-      // Set cluster ID in ClusterIdHolder for Schema Registry instrumentation
       KafkaConsumerInfo kafkaConsumerInfo =
           InstrumentationContext.get(KafkaConsumer.class, KafkaConsumerInfo.class).get(consumer);
+
+      // a poll() means any span deferred past the previous poll loop is done -- close it now (no-op
+      // unless consumer-scope deferral is enabled)
+      KafkaConsumerInstrumentationHelper.closeLingeringConsumeScope(kafkaConsumerInfo);
+
+      // Set cluster ID in ClusterIdHolder for Schema Registry instrumentation
       if (kafkaConsumerInfo != null && Config.get().isDataStreamsEnabled()) {
         Metadata consumerMetadata = kafkaConsumerInfo.getClientMetadata();
         if (consumerMetadata != null) {
@@ -290,4 +305,19 @@ public static void captureGroup(
       span.finish();
     }
   }
+
+  /**
+   * Finishes a {@code kafka.consume} span left active past the last poll when {@code
+   * close()}/{@code unsubscribe()} ends the poll loop, in case no further poll() arrives.
+   */
+  public static class CloseScopeAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(@Advice.This KafkaConsumer consumer) {
+      if (Config.get().isKafkaCreateConsumerScopeEnabled()) {
+        KafkaConsumerInfo kafkaConsumerInfo =
+            InstrumentationContext.get(KafkaConsumer.class, KafkaConsumerInfo.class).get(consumer);
+        KafkaConsumerInstrumentationHelper.closeLingeringConsumeScope(kafkaConsumerInfo);
+      }
+    }
+  }
 }

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

@@ -123,7 +123,13 @@ public static void wrap(
             KafkaConsumerInstrumentationHelper.extractBootstrapServers(kafkaConsumerInfo);
         iterable =
             new TracingIterable(
-                iterable, KAFKA_CONSUME, CONSUMER_DECORATE, group, clusterId, bootstrapServers);
+                iterable,
+                KAFKA_CONSUME,
+                CONSUMER_DECORATE,
+                group,
+                clusterId,
+                bootstrapServers,
+                kafkaConsumerInfo);
       }
     }
   }
@@ -145,7 +151,13 @@ public static void wrap(
             KafkaConsumerInstrumentationHelper.extractBootstrapServers(kafkaConsumerInfo);
         iterable =
             new TracingList(
-                iterable, KAFKA_CONSUME, CONSUMER_DECORATE, group, clusterId, bootstrapServers);
+                iterable,
+                KAFKA_CONSUME,
+                CONSUMER_DECORATE,
+                group,
+                clusterId,
+                bootstrapServers,
+                kafkaConsumerInfo);
       }
     }
   }
@@ -167,7 +179,13 @@ public static void wrap(
             KafkaConsumerInstrumentationHelper.extractBootstrapServers(kafkaConsumerInfo);
         iterator =
             new TracingIterator(
-                iterator, KAFKA_CONSUME, CONSUMER_DECORATE, group, clusterId, bootstrapServers);
+                iterator,
+                KAFKA_CONSUME,
+                CONSUMER_DECORATE,
+                group,
+                clusterId,
+                bootstrapServers,
+                kafkaConsumerInfo);
       }
     }
   }

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

@@ -1,10 +1,52 @@
 package datadog.trace.instrumentation.kafka_clients;
 
+import datadog.context.Context;
+import datadog.trace.api.Config;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.instrumentation.kafka_common.MetadataState;
 import org.apache.kafka.clients.Metadata;
 
 public class KafkaConsumerInstrumentationHelper {
+
+  /** Whether the last record's consume span should be kept active past the poll loop. */
+  public static boolean shouldDeferConsumerScope() {
+    return Config.get().isKafkaCreateConsumerScopeEnabled();
+  }
+
+  /**
+   * Finishes the {@code kafka.consume} span deliberately left active past the poll loop when
+   * consumer-scope deferral is enabled, restoring the caller's context. Safe to call from any
+   * consumer entry point (poll, close, unsubscribe).
+   */
+  public static void closeLingeringConsumeScope(KafkaConsumerInfo info) {
+    if (!shouldDeferConsumerScope()) {
+      return;
+    }
+    // restore the caller's context when the lingering consume span is still active on this thread
+    AgentSpan active = AgentTracer.activeSpan();
+    if (active != null && KafkaDecorator.KAFKA_CONSUME.equals(active.getOperationName())) {
+      if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
+        AgentTracer.closeLingeringIterationScope();
+      } else {
+        AgentSpan swappedOut = AgentSpan.fromContext(Context.root().swap());
+        if (swappedOut != null) {
+          swappedOut.finishWithEndToEnd();
+        }
+      }
+    }
+    // owner-aware: finish the exact deferred span regardless of thread/context (CAS-safe if already
+    // finished)
+    if (info != null) {
+      AgentSpan deferred = info.getAndClearDeferredConsumeSpan();
+      if (deferred != null) {
+        deferred.finishWithEndToEnd();
+      }
+    }
+  }
+
   public static String extractGroup(KafkaConsumerInfo kafkaConsumerInfo) {
     if (kafkaConsumerInfo != null) {
       return kafkaConsumerInfo.getConsumerGroup();

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

@@ -10,27 +10,36 @@ public class TracingIterable implements Iterable<ConsumerRecord<?, ?>>, TracingI
   private final String group;
   private final String clusterId;
   private final String bootstrapServers;
+  private final KafkaConsumerInfo kafkaConsumerInfo;
 
   public TracingIterable(
       final Iterable<ConsumerRecord<?, ?>> delegate,
       final CharSequence operationName,
       final KafkaDecorator decorator,
       String group,
       String clusterId,
-      String bootstrapServers) {
+      String bootstrapServers,
+      KafkaConsumerInfo kafkaConsumerInfo) {
     this.delegate = delegate;
     this.operationName = operationName;
     this.decorator = decorator;
     this.group = group;
     this.clusterId = clusterId;
     this.bootstrapServers = bootstrapServers;
+    this.kafkaConsumerInfo = kafkaConsumerInfo;
   }
 
   @Override
   public Iterator<ConsumerRecord<?, ?>> iterator() {
     // every iteration will add spans. Not only the very first one
     return new TracingIterator(
-        delegate.iterator(), operationName, decorator, group, clusterId, bootstrapServers);
+        delegate.iterator(),
+        operationName,
+        decorator,
+        group,
+        clusterId,
+        bootstrapServers,
+        kafkaConsumerInfo);
   }
 
   @Override

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

@@ -47,39 +47,61 @@ public class TracingIterator implements Iterator<ConsumerRecord<?, ?>> {
   private final String group;
   private final String clusterId;
   private final String bootstrapServers;
+  private final KafkaConsumerInfo kafkaConsumerInfo;
 
   public TracingIterator(
       final Iterator<ConsumerRecord<?, ?>> delegateIterator,
       final CharSequence operationName,
       final KafkaDecorator decorator,
       String group,
       String clusterId,
-      String bootstrapServers) {
+      String bootstrapServers,
+      KafkaConsumerInfo kafkaConsumerInfo) {
     this.delegateIterator = delegateIterator;
     this.operationName = operationName;
     this.decorator = decorator;
     this.group = group;
     this.clusterId = clusterId;
     this.bootstrapServers = bootstrapServers;
+    this.kafkaConsumerInfo = kafkaConsumerInfo;
   }
 
   @Override
   public boolean hasNext() {
     boolean moreRecords = delegateIterator.hasNext();
     if (!moreRecords) {
-      // no more records, use this as a signal to close the last iteration scope
-      if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
-        closePrevious(true);
-      } else {
-        final AgentSpan previousSpan = AgentSpan.fromContext(Context.root().swap());
-        if (previousSpan != null) {
-          previousSpan.finishWithEndToEnd();
+      if (!KafkaConsumerInstrumentationHelper.shouldDeferConsumerScope()) {
+        // no more records, use this as a signal to close the last iteration scope
+        if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
+          closePrevious(true);
+        } else {
+          final AgentSpan previousSpan = AgentSpan.fromContext(Context.root().swap());
+          if (previousSpan != null) {
+            previousSpan.finishWithEndToEnd();
+          }
         }
+      } else {
+        // deferring: leave the last record's scope active past the loop; record a handle to it
+        captureDeferredConsumeSpan();
       }
     }
     return moreRecords;
   }
 
+  /**
+   * Stashes the active {@code kafka.consume} span on the per-consumer {@link KafkaConsumerInfo} so
+   * a later poll()/close()/unsubscribe() can finish that precise span regardless of thread or
+   * stack.
+   */
+  protected void captureDeferredConsumeSpan() {
+    if (kafkaConsumerInfo != null) {
+      AgentSpan span = AgentTracer.activeSpan();
+      if (span != null && KafkaDecorator.KAFKA_CONSUME.equals(span.getOperationName())) {
+        kafkaConsumerInfo.setDeferredConsumeSpan(span);
+      }
+    }
+  }
+
   @Override
   public ConsumerRecord<?, ?> next() {
     final ConsumerRecord<?, ?> next = delegateIterator.next();
@@ -90,7 +112,15 @@ public boolean hasNext() {
   protected void startNewRecordSpan(ConsumerRecord<?, ?> val) {
     try {
       if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
-        closePrevious(true);
+        // real records always close the previous scope; the terminal close is deferred when the
+        // flag is on -- see hasNext().
+        if (val != null || !KafkaConsumerInstrumentationHelper.shouldDeferConsumerScope()) {
+          closePrevious(true);
+        } else {
+          // terminal null record while deferring: leave the last consume span active, record a
+          // handle so a later poll()/close()/unsubscribe() finishes it
+          captureDeferredConsumeSpan();
+        }
       } else if (val == null) { // previous message span was the last
         final AgentSpan previousSpan = AgentSpan.fromContext(Context.root().swap());
         if (previousSpan != null) {

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

@@ -14,20 +14,23 @@ public class TracingList implements List<ConsumerRecord<?, ?>>, TracingIterableD
   private final String group;
   private final String clusterId;
   private final String bootstrapServers;
+  private final KafkaConsumerInfo kafkaConsumerInfo;
 
   public TracingList(
       final List<ConsumerRecord<?, ?>> delegate,
       final CharSequence operationName,
       final KafkaDecorator decorator,
       String group,
       String clusterId,
-      String bootstrapServers) {
+      String bootstrapServers,
+      KafkaConsumerInfo kafkaConsumerInfo) {
     this.operationName = operationName;
     this.decorator = decorator;
     this.delegate = delegate;
     this.group = group;
     this.clusterId = clusterId;
     this.bootstrapServers = bootstrapServers;
+    this.kafkaConsumerInfo = kafkaConsumerInfo;
   }
 
   @Override
@@ -140,7 +143,13 @@ public int lastIndexOf(final Object o) {
   public ListIterator<ConsumerRecord<?, ?>> listIterator(final int index) {
     // every iteration will add spans. Not only the very first one
     return new TracingListIterator(
-        delegate.listIterator(index), operationName, decorator, group, clusterId, bootstrapServers);
+        delegate.listIterator(index),
+        operationName,
+        decorator,
+        group,
+        clusterId,
+        bootstrapServers,
+        kafkaConsumerInfo);
   }
 
   @Override
@@ -151,7 +160,8 @@ public int lastIndexOf(final Object o) {
         decorator,
         group,
         clusterId,
-        bootstrapServers);
+        bootstrapServers,
+        kafkaConsumerInfo);
   }
 
   @Override