feat(dsm): Migrate DSM injection call to propagator API 3/3

Author: PerfectSlayer

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

@@ -3,9 +3,10 @@
 import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.api.datastreams.DataStreamsContext.fromTagsWithoutCheckpoint;
+import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_OUT;
 import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG;
@@ -24,9 +25,12 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
+import datadog.context.propagation.Propagator;
+import datadog.context.propagation.Propagators;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.Config;
+import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -155,8 +159,9 @@ public static AgentScope onEnter(
           // inject the context in the headers, but delay sending the stats until we know the
           // message size.
           // The stats are saved in the pathway context and sent in PayloadSizeAdvice.
-          propagate()
-              .injectPathwayContextWithoutSendingStats(span, record.headers(), setter, sortedTags);
+          Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN);
+          DataStreamsContext dsmContext = fromTagsWithoutCheckpoint(sortedTags);
+          dsmPropagator.inject(span.with(dsmContext), record.headers(), setter);
           AvroSchemaExtractor.tryExtractProducer(record, span);
         }
       } catch (final IllegalStateException e) {
@@ -173,8 +178,9 @@ record =
         defaultPropagator().inject(span, record.headers(), setter);
         if (STREAMING_CONTEXT.isDisabledForTopic(record.topic())
             || STREAMING_CONTEXT.isSinkTopic(record.topic())) {
-          propagate()
-              .injectPathwayContextWithoutSendingStats(span, record.headers(), setter, sortedTags);
+          Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN);
+          DataStreamsContext dsmContext = fromTagsWithoutCheckpoint(sortedTags);
+          dsmPropagator.inject(span.with(dsmContext), record.headers(), setter);
           AvroSchemaExtractor.tryExtractProducer(record, span);
         }
       }

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

@@ -1,9 +1,10 @@
 package datadog.trace.instrumentation.kafka_clients38;
 
 import static datadog.context.propagation.Propagators.defaultPropagator;
+import static datadog.trace.api.datastreams.DataStreamsContext.fromTagsWithoutCheckpoint;
+import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_OUT;
 import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG;
@@ -15,7 +16,10 @@
 import static datadog.trace.instrumentation.kafka_clients38.KafkaDecorator.TIME_IN_QUEUE_ENABLED;
 import static datadog.trace.instrumentation.kafka_common.StreamingContext.STREAMING_CONTEXT;
 
+import datadog.context.propagation.Propagator;
+import datadog.context.propagation.Propagators;
 import datadog.trace.api.Config;
+import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -79,8 +83,9 @@ public static AgentScope onEnter(
         // inject the context in the headers, but delay sending the stats until we know the
         // message size.
         // The stats are saved in the pathway context and sent in PayloadSizeAdvice.
-        propagate()
-            .injectPathwayContextWithoutSendingStats(span, record.headers(), setter, sortedTags);
+        Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN);
+        DataStreamsContext dsmContext = fromTagsWithoutCheckpoint(sortedTags);
+        dsmPropagator.inject(span.with(dsmContext), record.headers(), setter);
         AvroSchemaExtractor.tryExtractProducer(record, span);
       }
     } catch (final IllegalStateException e) {
@@ -97,8 +102,9 @@ record =
       defaultPropagator().inject(span, record.headers(), setter);
       if (STREAMING_CONTEXT.isDisabledForTopic(record.topic())
           || STREAMING_CONTEXT.isSinkTopic(record.topic())) {
-        propagate()
-            .injectPathwayContextWithoutSendingStats(span, record.headers(), setter, sortedTags);
+        Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN);
+        DataStreamsContext dsmContext = fromTagsWithoutCheckpoint(sortedTags);
+        dsmPropagator.inject(span.with(dsmContext), record.headers(), setter);
         AvroSchemaExtractor.tryExtractProducer(record, span);
       }
     }