Skip to content

Migrate context injection calls to context-first APIs #8358

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Feb 10, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .circleci/config.continue.yml.j2
Original file line number Diff line number Diff line change
Expand Up @@ -767,7 +767,7 @@ jobs:

muzzle:
<<: *defaults
resource_class: medium
resource_class: medium+
parallelism: 4
steps:
- setup_code
Expand Down Expand Up @@ -798,7 +798,7 @@ jobs:
name: Verify Muzzle
command: >-
SKIP_BUILDSCAN="true"
GRADLE_OPTS="-Dorg.gradle.jvmargs='-Xmx2G -Xms2G -XX:ErrorFile=/tmp/hs_err_pid%p.log -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/tmp'"
GRADLE_OPTS="-Dorg.gradle.jvmargs='-Xmx3G -Xms2G -XX:ErrorFile=/tmp/hs_err_pid%p.log -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/tmp'"
./gradlew `circleci tests split --split-by=timings workspace/build/muzzleTasks | xargs`
<< pipeline.parameters.gradle_flags >>
--max-workers=4
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,9 @@

import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
import java.net.HttpURLConnection;
import javax.annotation.ParametersAreNonnullByDefault;

@ParametersAreNonnullByDefault
public class HeadersInjectAdapter implements AgentPropagation.Setter<HttpURLConnection> {

public static final HeadersInjectAdapter SETTER = new HeadersInjectAdapter();
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
package datadog.trace.bootstrap.instrumentation.rmi;

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.context.propagation.Propagators.defaultPropagator;

import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
Expand All @@ -9,6 +9,7 @@
import java.io.ObjectOutput;
import java.util.HashMap;
import java.util.Map;
import javax.annotation.ParametersAreNonnullByDefault;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

Expand All @@ -33,7 +34,7 @@ public Map<String, String> getContext() {

public static ContextPayload from(final AgentSpan span) {
final ContextPayload payload = new ContextPayload();
propagate().inject(span, payload, SETTER);
defaultPropagator().inject(span, payload, SETTER);
return payload;
}

Expand All @@ -54,6 +55,7 @@ public void write(final ObjectOutput out) throws IOException {
out.writeObject(context);
}

@ParametersAreNonnullByDefault
public static class InjectAdapter implements AgentPropagation.Setter<ContextPayload> {
@Override
public void set(final ContextPayload carrier, final String key, final String value) {
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
package datadog.trace.civisibility.domain.buildsystem;

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.context.propagation.Propagators.defaultPropagator;

import datadog.communication.ddagent.TracerVersion;
import datadog.trace.api.Config;
Expand Down Expand Up @@ -33,10 +33,15 @@
import datadog.trace.util.Strings;
import java.net.InetSocketAddress;
import java.nio.file.Path;
import java.util.*;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.atomic.LongAdder;
import java.util.function.Consumer;
import javax.annotation.Nullable;
import javax.annotation.ParametersAreNonnullByDefault;

public class BuildSystemModuleImpl extends AbstractTestModule implements BuildSystemModule {

Expand Down Expand Up @@ -105,6 +110,7 @@ public <T extends CoverageCalculator> BuildSystemModuleImpl(
setTag(Tags.TEST_COMMAND, startCommand);
}

@ParametersAreNonnullByDefault
private static final class ChildProcessPropertiesPropagationSetter
implements AgentPropagation.Setter<Map<String, String>> {
static final AgentPropagation.Setter<Map<String, String>> INSTANCE =
Expand Down Expand Up @@ -221,7 +227,7 @@ private Map<String, String> getPropertiesPropagatedToChildProcess(
}

// propagate module span context to child processes
propagate()
defaultPropagator()
.inject(span, propagatedSystemProperties, ChildProcessPropertiesPropagationSetter.INSTANCE);

return propagatedSystemProperties;
Expand Down
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
package datadog.opentelemetry.shim.context.propagation;

import static datadog.context.propagation.Propagators.defaultPropagator;
import static datadog.opentelemetry.shim.trace.OtelSpanContext.fromRemote;
import static datadog.trace.api.TracePropagationStyle.TRACECONTEXT;

import datadog.opentelemetry.shim.context.OtelContext;
import datadog.opentelemetry.shim.trace.OtelExtractedContext;
import datadog.opentelemetry.shim.trace.OtelSpan;
import datadog.trace.api.TracePropagationStyle;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext.Extracted;
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
Expand Down Expand Up @@ -36,11 +38,7 @@ public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> se
if (carrier == null) {
return;
}
Span span = Span.fromContext(context);
if (span.getSpanContext().isValid()) {
AgentSpanContext agentSpanContext = OtelExtractedContext.extract(context);
AgentTracer.propagate().inject(agentSpanContext, carrier, setter::set);
}
defaultPropagator().inject(convertContext(context), carrier, setter::set);
}

@Override
Expand All @@ -66,15 +64,23 @@ public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C
}
}

private static datadog.context.Context convertContext(Context context) {
// TODO Extract baggage too
// TODO Create fast path from OtelSpan --> AgentSpan delegate --> with() to inflate as full
// context if baggage
AgentSpanContext extract = OtelExtractedContext.extract(context);
return AgentSpan.fromSpanContext(extract);
}

/**
* Extracts tracestate if {@code tracestate} header is present and extracted context comes from
* {@link TracePropagationStyle#TRACECONTEXT}
*
* @param extracted The extracted context.
* @param carrier The context carrier.
* @param getter The context getter.
* @return The extracted tracestate, or an empty tracestate otherwise.
* @param <C> The carrier type.
* @return The extracted tracestate, or an empty tracestate otherwise.
*/
private static <C> TraceState extractTraceState(
Extracted extracted, C carrier, TextMapGetter<C> getter) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,7 @@ public AgentSpanContext getAgentSpanContext() {

@Override
public AgentSpan asAgentSpan() {
return delegate;
return this.delegate;
}

private static class NoopSpan implements Span {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import akka.http.scaladsl.model.headers.CustomHeader;
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import javax.annotation.ParametersAreNonnullByDefault;
import scala.runtime.AbstractFunction1;
import scala.util.Try;

Expand Down Expand Up @@ -35,7 +36,7 @@ public Void apply(final Try<HttpResponse> result) {
public static class AkkaHttpHeaders implements AgentPropagation.Setter<HttpRequest> {
private HttpRequest request;
// Did this request have a span when the AkkaHttpHeaders object was created?
private boolean hadSpan;
private final boolean hadSpan;

public AkkaHttpHeaders(final HttpRequest request) {
hadSpan = request != null && request.getHeader(HasSpanHeader.class).isPresent();
Expand All @@ -51,6 +52,7 @@ public boolean hadSpan() {
return hadSpan;
}

@ParametersAreNonnullByDefault
@Override
public void set(final HttpRequest carrier, final String key, final String value) {
// Coerce a Scala trait Self type into the correct type
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package datadog.trace.instrumentation.akkahttp;

import static datadog.context.propagation.Propagators.defaultPropagator;
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
Expand Down Expand Up @@ -78,7 +79,7 @@ public static AgentScope methodEnter(
DECORATE.onRequest(span, request);

if (request != null) {
propagate().inject(span, request, headers);
defaultPropagator().inject(span, request, headers);
propagate()
.injectPathwayContext(
span, request, headers, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package datadog.trace.instrumentation.akkahttp106;

import static datadog.context.propagation.Propagators.defaultPropagator;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
Expand Down Expand Up @@ -28,7 +29,7 @@ public static AgentScope methodEnter(
AkkaHttpClientDecorator.DECORATE.onRequest(span, request);

if (request != null) {
propagate().inject(span, request, headers);
defaultPropagator().inject(span, request, headers);
propagate()
.injectPathwayContext(
span, request, headers, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package datadog.trace.instrumentation.apachehttpasyncclient;

import static datadog.context.propagation.Propagators.defaultPropagator;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
import static datadog.trace.instrumentation.apachehttpasyncclient.HttpHeadersInjectAdapter.SETTER;
Expand Down Expand Up @@ -34,7 +35,7 @@ public HttpRequest generateRequest() throws IOException, HttpException {
final HttpRequest request = delegate.generateRequest();
DECORATE.onRequest(span, new HostAndRequestAsHttpUriRequest(delegate.getTarget(), request));

propagate().inject(span, request, SETTER);
defaultPropagator().inject(span, request, SETTER);
propagate()
.injectPathwayContext(span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);

Expand Down
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
package datadog.trace.instrumentation.apachehttpasyncclient;

import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
import javax.annotation.ParametersAreNonnullByDefault;
import org.apache.http.HttpRequest;

@ParametersAreNonnullByDefault
public class HttpHeadersInjectAdapter implements AgentPropagation.Setter<HttpRequest> {

public static final HttpHeadersInjectAdapter SETTER = new HttpHeadersInjectAdapter();
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package datadog.trace.instrumentation.apachehttpclient;

import static datadog.context.propagation.Propagators.defaultPropagator;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
Expand Down Expand Up @@ -45,7 +46,7 @@ private static AgentScope activateHttpSpan(final HttpUriRequest request) {

// AWS calls are often signed, so we can't add headers without breaking the signature.
if (!awsClientCall) {
propagate().inject(span, request, SETTER);
defaultPropagator().inject(span, request, SETTER);
propagate()
.injectPathwayContext(
span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
Expand Down
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
package datadog.trace.instrumentation.apachehttpclient;

import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
import javax.annotation.ParametersAreNonnullByDefault;
import org.apache.http.client.methods.HttpUriRequest;

@ParametersAreNonnullByDefault
public class HttpHeadersInjectAdapter implements AgentPropagation.Setter<HttpUriRequest> {

public static final HttpHeadersInjectAdapter SETTER = new HttpHeadersInjectAdapter();
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package datadog.trace.instrumentation.apachehttpclient5;

import static datadog.context.propagation.Propagators.defaultPropagator;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
Expand Down Expand Up @@ -27,7 +28,7 @@ public void sendRequest(HttpRequest request, EntityDetails entityDetails, HttpCo
throws HttpException, IOException {
DECORATE.onRequest(span, request);

propagate().inject(span, request, SETTER);
defaultPropagator().inject(span, request, SETTER);
propagate()
.injectPathwayContext(span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
delegate.sendRequest(request, entityDetails, context);
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package datadog.trace.instrumentation.apachehttpclient5;

import static datadog.context.propagation.Propagators.defaultPropagator;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
Expand Down Expand Up @@ -45,7 +46,7 @@ private static AgentScope activateHttpSpan(final HttpRequest request) {
final boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
// AWS calls are often signed, so we can't add headers without breaking the signature.
if (!awsClientCall) {
propagate().inject(span, request, SETTER);
defaultPropagator().inject(span, request, SETTER);
propagate()
.injectPathwayContext(
span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
Expand Down
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
package datadog.trace.instrumentation.apachehttpclient5;

import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
import javax.annotation.ParametersAreNonnullByDefault;
import org.apache.hc.core5.http.HttpRequest;

@ParametersAreNonnullByDefault
public class HttpHeadersInjectAdapter implements AgentPropagation.Setter<HttpRequest> {

public static final HttpHeadersInjectAdapter SETTER = new HttpHeadersInjectAdapter();
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package datadog.trace.instrumentation.armeria.grpc.client;

import static datadog.context.propagation.Propagators.defaultPropagator;
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
Expand Down Expand Up @@ -120,7 +121,7 @@ public static <T> AgentScope before(
if (null != responseListener && null != headers) {
span = InstrumentationContext.get(ClientCall.class, AgentSpan.class).get(call);
if (null != span) {
propagate().inject(span, headers, SETTER);
defaultPropagator().inject(span, headers, SETTER);
propagate().injectPathwayContext(span, headers, SETTER, CLIENT_PATHWAY_EDGE_TAGS);
return activateSpan(span);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,9 @@

import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
import io.grpc.Metadata;
import javax.annotation.ParametersAreNonnullByDefault;

@ParametersAreNonnullByDefault
public final class GrpcInjectAdapter implements AgentPropagation.Setter<Metadata> {

public static final GrpcInjectAdapter SETTER = new GrpcInjectAdapter();
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package datadog.trace.instrumentation.aws.v2.eventbridge;

import static datadog.context.propagation.Propagators.defaultPropagator;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
import static datadog.trace.core.datastreams.TagsProcessor.BUS_TAG;
Expand Down Expand Up @@ -86,7 +87,7 @@ private String getTraceContextToInject(
jsonBuilder.append('{');

// Inject trace context
propagate().inject(span, jsonBuilder, SETTER);
defaultPropagator().inject(span, jsonBuilder, SETTER);

if (traceConfig().isDataStreamsEnabled()) {
propagate().injectPathwayContext(span, jsonBuilder, SETTER, getTags(eventBusName));
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
package datadog.trace.instrumentation.aws.v2.eventbridge;

import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
import javax.annotation.ParametersAreNonnullByDefault;

@ParametersAreNonnullByDefault
public class TextMapInjectAdapter implements AgentPropagation.Setter<StringBuilder> {

public static final TextMapInjectAdapter SETTER = new TextMapInjectAdapter();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import java.util.Locale;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.annotation.ParametersAreNonnullByDefault;

public class AwsSdkClientDecorator extends HttpClientDecorator<Request, Response>
implements AgentPropagation.Setter<Request<?>> {
Expand Down Expand Up @@ -332,6 +333,7 @@ protected int status(final Response response) {
return response.getHttpResponse().getStatusCode();
}

@ParametersAreNonnullByDefault
@Override
public void set(Request<?> carrier, String key, String value) {
carrier.addHeader(key, value);
Expand Down
Loading
Loading