Skip to content

Migrate instrumentations to Context API #9358

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

Open
wants to merge 8 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 4 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
Original file line number Diff line number Diff line change
@@ -1,14 +1,15 @@
package datadog.trace.bootstrap.instrumentation.decorator;

import static datadog.context.Context.root;
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
import static java.util.Collections.emptyMap;
import static java.util.concurrent.TimeUnit.MICROSECONDS;
import static java.util.concurrent.TimeUnit.SECONDS;

import datadog.context.Context;
import datadog.trace.api.GlobalTracer;
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
import datadog.trace.bootstrap.instrumentation.api.ContextVisitors;
import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
import datadog.trace.bootstrap.instrumentation.api.URIDefaultDataAdapter;
Expand All @@ -17,7 +18,6 @@
import datadog.trace.core.CoreTracer;
import datadog.trace.core.DDSpan;
import java.net.URI;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import org.openjdk.jmh.annotations.Benchmark;
Expand Down Expand Up @@ -59,13 +59,13 @@ public void setUp() {
.build();
GlobalTracer.forceRegister(tracer);
decorator = new BenchmarkHttpServerDecorator();
Context context = decorator.startSpan(Collections.emptyMap(), Context.root());
Context context = decorator.startSpan(emptyMap(), root());
span = fromContext(context);
}

@Benchmark
public AgentSpan onRequest() {
return decorator.onRequest(span, null, request, (AgentSpanContext.Extracted) null);
return decorator.onRequest(span, null, request, root());
}

public static class Request {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import static datadog.trace.api.cache.RadixTreeCache.UNSET_STATUS;
import static datadog.trace.api.datastreams.DataStreamsContext.forHttpServer;
import static datadog.trace.api.gateway.Events.EVENTS;
import static datadog.trace.bootstrap.ActiveSubsystems.APPSEC_ACTIVE;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;

Expand All @@ -17,11 +18,11 @@
import datadog.trace.api.gateway.BlockResponseFunction;
import datadog.trace.api.gateway.CallbackProvider;
import datadog.trace.api.gateway.Flow;
import datadog.trace.api.gateway.Flow.Action.RequestBlockingAction;
import datadog.trace.api.gateway.IGSpanInfo;
import datadog.trace.api.gateway.RequestContext;
import datadog.trace.api.gateway.RequestContextSlot;
import datadog.trace.api.naming.SpanNaming;
import datadog.trace.bootstrap.ActiveSubsystems;
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
Expand All @@ -44,7 +45,6 @@
import java.util.function.Function;
import java.util.function.Supplier;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

Expand Down Expand Up @@ -150,8 +150,8 @@ public Context startSpan(REQUEST_CARRIER carrier, Context context) {
AgentSpan span =
tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true);
Flow<Void> flow = callIGCallbackRequestHeaders(span, carrier);
if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
if (flow.getAction() instanceof RequestBlockingAction) {
span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
}
AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
if (null != carrier && null != getter) {
Expand All @@ -165,25 +165,9 @@ public AgentSpan onRequest(
final CONNECTION connection,
final REQUEST request,
final Context context) {
return onRequest(span, connection, request, getExtractedSpanContext(context));
}

public AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
AgentSpan extractedSpan = AgentSpan.fromContext(context);
return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
}

public AgentSpan onRequest(
final AgentSpan span,
final CONNECTION connection,
final REQUEST request,
@Nullable final AgentSpanContext.Extracted context) {
Config config = Config.get();
boolean clientIpResolverEnabled =
config.isClientIpEnabled()
|| traceClientIpResolverEnabled && ActiveSubsystems.APPSEC_ACTIVE;

if (ActiveSubsystems.APPSEC_ACTIVE) {
if (APPSEC_ACTIVE) {
RequestContext requestContext = span.getRequestContext();
if (requestContext != null) {
BlockResponseFunction brf = createBlockResponseFunction(request, connection);
Expand All @@ -193,35 +177,38 @@ public AgentSpan onRequest(
}
Flow<Void> flow = callIGCallbackRequestSessionId(span, request);
Flow.Action action = flow.getAction();
if (action instanceof Flow.Action.RequestBlockingAction) {
span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
if (action instanceof RequestBlockingAction) {
span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
}
}

if (context != null) {
AgentSpanContext.Extracted extracted = callIGCallbackStart(context);
boolean clientIpResolverEnabled =
config.isClientIpEnabled() || traceClientIpResolverEnabled && APPSEC_ACTIVE;
if (extracted != null) {
if (clientIpResolverEnabled) {
String forwarded = context.getForwarded();
String forwarded = extracted.getForwarded();
if (forwarded != null) {
span.setTag(Tags.HTTP_FORWARDED, forwarded);
}
String forwardedProto = context.getXForwardedProto();
String forwardedProto = extracted.getXForwardedProto();
if (forwardedProto != null) {
span.setTag(Tags.HTTP_FORWARDED_PROTO, forwardedProto);
}
String forwardedHost = context.getXForwardedHost();
String forwardedHost = extracted.getXForwardedHost();
if (forwardedHost != null) {
span.setTag(Tags.HTTP_FORWARDED_HOST, forwardedHost);
}
String forwardedIp = context.getXForwardedFor();
String forwardedIp = extracted.getXForwardedFor();
if (forwardedIp != null) {
span.setTag(Tags.HTTP_FORWARDED_IP, forwardedIp);
}
String forwardedPort = context.getXForwardedPort();
String forwardedPort = extracted.getXForwardedPort();
if (forwardedPort != null) {
span.setTag(Tags.HTTP_FORWARDED_PORT, forwardedPort);
}
}
String userAgent = context.getUserAgent();
String userAgent = extracted.getUserAgent();
if (userAgent != null) {
span.setTag(Tags.HTTP_USER_AGENT, userAgent);
}
Expand All @@ -245,8 +232,8 @@ public AgentSpan onRequest(
} else if (supportsRaw) {
span.setTag(Tags.HTTP_URL, URIUtils.lazyInvalidUrl(url.raw()));
}
if (context != null && context.getXForwardedHost() != null) {
span.setTag(Tags.HTTP_HOSTNAME, context.getXForwardedHost());
if (extracted != null && extracted.getXForwardedHost() != null) {
span.setTag(Tags.HTTP_HOSTNAME, extracted.getXForwardedHost());
} else if (url.host() != null) {
span.setTag(Tags.HTTP_HOSTNAME, url.host());
}
Expand All @@ -258,8 +245,8 @@ public AgentSpan onRequest(
span.setTag(DDTags.HTTP_FRAGMENT, url.fragment());
}
Flow<Void> flow = callIGCallbackURI(span, url, method);
if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
if (flow.getAction() instanceof RequestBlockingAction) {
span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
}
if (valid && SHOULD_SET_URL_RESOURCE_NAME) {
HTTP_RESOURCE_DECORATOR.withServerPath(span, method, path, encoded);
Expand All @@ -280,8 +267,8 @@ public AgentSpan onRequest(
}

String inferredAddressStr = null;
if (clientIpResolverEnabled && context != null) {
InetAddress inferredAddress = ClientIpAddressResolver.resolve(context, span);
if (clientIpResolverEnabled && extracted != null) {
InetAddress inferredAddress = ClientIpAddressResolver.resolve(extracted, span);
// the peer address should be used if:
// 1. the headers yield nothing, regardless of whether it is public or not
// 2. it is public and the headers yield a private address
Expand All @@ -300,9 +287,9 @@ public AgentSpan onRequest(
span.setTag(Tags.HTTP_CLIENT_IP, inferredAddressStr);
}
} else if (clientIpResolverEnabled && span.getLocalRootSpan() != span) {
// in this case context == null
// If there is no context we can't do anything but use the peer addr.
// Additionally, context == null arises on subspans for which the resolution
// in this case extracted == null
// If there is no extracted we can't do anything but use the peer addr.
// Additionally, extracted == null arises on subspans for which the resolution
// likely already happened on the top span, so we don't need to do the resolution
// again. Instead, copy from the top span, should it exist
AgentSpan localRootSpan = span.getLocalRootSpan();
Expand All @@ -321,13 +308,18 @@ public AgentSpan onRequest(
}
setPeerPort(span, peerPort);
Flow<Void> flow = callIGCallbackAddressAndPort(span, peerIp, peerPort, inferredAddressStr);
if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
if (flow.getAction() instanceof RequestBlockingAction) {
span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
}

return span;
}

protected static AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
AgentSpan extractedSpan = AgentSpan.fromContext(context);
return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
}

protected BlockResponseFunction createBlockResponseFunction(
REQUEST request, CONNECTION connection) {
return null;
Expand Down
Loading
Loading