Skip to content

Commit ad25479

Browse files
PerfectSlayerzarirhamza
PerfectSlayer
and
zarirhamza
authored
Refactor Severless Gateway Inferred Span (#9388)
* Implement APIGW Inferred Proxy Spans (#8336) * Creates inferred proxy spans for API Gateway calls via presence of http headers Co-authored-by: Zarir Hamza <[email protected]> * feat(serverless): Fix gateway inferred span design Avoid duplicate expensive context extraction Avoid subclassing tracing span for serverless but used serverless context element instead to store / track inferred span while keep tracing feature untouched Improved propagator to not create / capture inferred span context element on invalid data Rework context element to hold the inferred spans and its captured data Release captured data as soon as they span start (never read after this point so reclaiming memory) Refactor context element and propagator into the right package, not context component (product / feature agnostic) Refactor unit tests --------- Co-authored-by: Zarir Hamza <[email protected]>
1 parent 4313a72 commit ad25479

File tree

10 files changed

+407
-11
lines changed

10 files changed

+407
-11
lines changed

‎.github/CODEOWNERS‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,10 @@
2828
/internal-api/src/test/groovy/datadog/trace/api/sampling @DataDog/apm-sdk-api-java
2929

3030
# @DataDog/apm-serverless
31-
/dd-trace-core/src/main/java/datadog/trace/lambda/ @DataDog/apm-serverless
32-
/dd-trace-core/src/test/groovy/datadog/trace/lambda/ @DataDog/apm-serverless
31+
/dd-trace-core/src/main/java/datadog/trace/lambda/ @DataDog/apm-serverless
32+
/dd-trace-core/src/test/groovy/datadog/trace/lambda/ @DataDog/apm-serverless
33+
**/InferredProxy*.java @DataDog/apm-serverless
34+
**/InferredProxy*.groovy @DataDog/apm-serverless
3335

3436
# @DataDog/apm-lang-platform-java
3537
/.circleci/ @DataDog/apm-lang-platform-java

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java‎

Lines changed: 29 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import datadog.trace.api.gateway.Flow;
2121
import datadog.trace.api.gateway.Flow.Action.RequestBlockingAction;
2222
import datadog.trace.api.gateway.IGSpanInfo;
23+
import datadog.trace.api.gateway.InferredProxySpan;
2324
import datadog.trace.api.gateway.RequestContext;
2425
import datadog.trace.api.gateway.RequestContextSlot;
2526
import datadog.trace.api.naming.SpanNaming;
@@ -147,21 +148,32 @@ public Context startSpan(REQUEST_CARRIER carrier, Context parentContext) {
147148
instrumentationNames != null && instrumentationNames.length > 0
148149
? instrumentationNames[0]
149150
: DEFAULT_INSTRUMENTATION_NAME;
150-
AgentSpanContext.Extracted extracted =
151-
callIGCallbackStart(getExtractedSpanContext(parentContext));
151+
AgentSpanContext extracted = getExtractedSpanContext(parentContext);
152+
// Call IG callbacks
153+
extracted = callIGCallbackStart(extracted);
154+
// Create gateway inferred span if needed
155+
extracted = startInferredProxySpan(parentContext, extracted);
152156
AgentSpan span =
153157
tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true);
158+
// Apply RequestBlockingAction if any
154159
Flow<Void> flow = callIGCallbackRequestHeaders(span, carrier);
155160
if (flow.getAction() instanceof RequestBlockingAction) {
156161
span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
157162
}
158-
AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
159-
if (null != carrier && null != getter) {
160-
tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer());
161-
}
163+
// DSM Checkpoint
164+
tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer());
162165
return parentContext.with(span);
163166
}
164167

168+
protected AgentSpanContext startInferredProxySpan(Context context, AgentSpanContext extracted) {
169+
InferredProxySpan span;
170+
if (!Config.get().isInferredProxyPropagationEnabled()
171+
|| (span = InferredProxySpan.fromContext(context)) == null) {
172+
return extracted;
173+
}
174+
return span.start(extracted);
175+
}
176+
165177
public AgentSpan onRequest(
166178
final AgentSpan span,
167179
final CONNECTION connection,
@@ -390,8 +402,7 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) {
390402
return span;
391403
}
392404

393-
private AgentSpanContext.Extracted callIGCallbackStart(
394-
@Nullable final AgentSpanContext.Extracted extracted) {
405+
private AgentSpanContext callIGCallbackStart(@Nullable final AgentSpanContext extracted) {
395406
AgentTracer.TracerAPI tracer = tracer();
396407
Supplier<Flow<Object>> startedCbAppSec =
397408
tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted());
@@ -527,10 +538,20 @@ private Flow<Void> callIGCallbackURI(
527538

528539
@Override
529540
public AgentSpan beforeFinish(AgentSpan span) {
541+
// TODO Migrate beforeFinish to Context API
530542
onRequestEndForInstrumentationGateway(span);
543+
// Close Serverless Gateway Inferred Span if any
544+
// finishInferredProxySpan(context);
531545
return super.beforeFinish(span);
532546
}
533547

548+
protected void finishInferredProxySpan(Context context) {
549+
InferredProxySpan span;
550+
if ((span = InferredProxySpan.fromContext(context)) != null) {
551+
span.finish();
552+
}
553+
}
554+
534555
private void onRequestEndForInstrumentationGateway(@Nonnull final AgentSpan span) {
535556
if (span.getLocalRootSpan() != span) {
536557
return;

‎dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,9 @@ public final class TracerConfig {
100100
public static final String TRACE_BAGGAGE_MAX_BYTES = "trace.baggage.max.bytes";
101101
public static final String TRACE_BAGGAGE_TAG_KEYS = "trace.baggage.tag.keys";
102102

103+
public static final String TRACE_INFERRED_PROXY_SERVICES_ENABLED =
104+
"trace.inferred.proxy.services.enabled";
105+
103106
public static final String ENABLE_TRACE_AGENT_V05 = "trace.agent.v0.5.enabled";
104107

105108
public static final String CLIENT_IP_ENABLED = "trace.client-ip.enabled";

‎dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import static datadog.trace.api.TracePropagationBehaviorExtract.IGNORE;
88
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.BAGGAGE_CONCERN;
99
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN;
10+
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.INFERRED_PROXY_CONCERN;
1011
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.TRACING_CONCERN;
1112
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
1213
import static datadog.trace.common.metrics.MetricsAggregatorFactory.createMetricsAggregator;
@@ -90,6 +91,7 @@
9091
import datadog.trace.core.monitor.TracerHealthMetrics;
9192
import datadog.trace.core.propagation.ExtractedContext;
9293
import datadog.trace.core.propagation.HttpCodec;
94+
import datadog.trace.core.propagation.InferredProxyPropagator;
9395
import datadog.trace.core.propagation.PropagationTags;
9496
import datadog.trace.core.propagation.TracingPropagator;
9597
import datadog.trace.core.propagation.XRayPropagator;
@@ -816,6 +818,9 @@ private CoreTracer(
816818
&& config.getTracePropagationBehaviorExtract() != IGNORE) {
817819
Propagators.register(BAGGAGE_CONCERN, new BaggagePropagator(config));
818820
}
821+
if (config.isInferredProxyPropagationEnabled()) {
822+
Propagators.register(INFERRED_PROXY_CONCERN, new InferredProxyPropagator());
823+
}
819824

820825
if (config.isCiVisibilityEnabled()) {
821826
if (config.isCiVisibilityTraceSanitationEnabled()) {

‎dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
package datadog.trace.core.propagation;
2+
3+
import static datadog.trace.api.gateway.InferredProxySpan.fromHeaders;
4+
5+
import datadog.context.Context;
6+
import datadog.context.propagation.CarrierSetter;
7+
import datadog.context.propagation.CarrierVisitor;
8+
import datadog.context.propagation.Propagator;
9+
import datadog.trace.api.gateway.InferredProxySpan;
10+
import java.util.HashMap;
11+
import java.util.Map;
12+
import java.util.function.BiConsumer;
13+
import javax.annotation.ParametersAreNonnullByDefault;
14+
15+
/** Inferred proxy propagator. Only extract, not meant for injection. */
16+
@ParametersAreNonnullByDefault
17+
public class InferredProxyPropagator implements Propagator {
18+
private static final String INFERRED_PROXY_KEY_PREFIX = "x-dd-proxy";
19+
20+
@Override
21+
public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {}
22+
23+
@Override
24+
public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
25+
if (context == null || carrier == null || visitor == null) {
26+
return context;
27+
}
28+
InferredProxyContextExtractor extractor = new InferredProxyContextExtractor();
29+
visitor.forEachKeyValue(carrier, extractor);
30+
InferredProxySpan inferredProxySpan = extractor.inferredProxySpan();
31+
if (inferredProxySpan != null) {
32+
context = context.with(inferredProxySpan);
33+
}
34+
return context;
35+
}
36+
37+
/** Extract inferred proxy related headers into a map. */
38+
private static class InferredProxyContextExtractor implements BiConsumer<String, String> {
39+
private Map<String, String> values;
40+
41+
@Override
42+
public void accept(String key, String value) {
43+
if (key == null || key.isEmpty() || !key.startsWith(INFERRED_PROXY_KEY_PREFIX)) {
44+
return;
45+
}
46+
if (values == null) {
47+
this.values = new HashMap<>();
48+
}
49+
this.values.put(key, value);
50+
}
51+
52+
public InferredProxySpan inferredProxySpan() {
53+
if (this.values == null) {
54+
return null;
55+
}
56+
InferredProxySpan inferredProxySpan = fromHeaders(this.values);
57+
return inferredProxySpan.isValid() ? inferredProxySpan : null;
58+
}
59+
}
60+
}

‎dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
package datadog.trace.core.propagation;
2+
3+
import static datadog.context.Context.root;
4+
import static datadog.trace.api.gateway.InferredProxySpan.fromContext;
5+
import static org.junit.jupiter.api.Assertions.assertNotNull;
6+
import static org.junit.jupiter.api.Assertions.assertNull;
7+
import static org.junit.jupiter.api.Assertions.assertTrue;
8+
import static org.junit.jupiter.params.provider.Arguments.of;
9+
10+
import datadog.context.Context;
11+
import datadog.context.propagation.CarrierVisitor;
12+
import datadog.trace.api.gateway.InferredProxySpan;
13+
import java.util.HashMap;
14+
import java.util.Map;
15+
import java.util.function.BiConsumer;
16+
import java.util.stream.Stream;
17+
import javax.annotation.ParametersAreNonnullByDefault;
18+
import org.junit.jupiter.api.BeforeEach;
19+
import org.junit.jupiter.api.DisplayName;
20+
import org.junit.jupiter.api.Test;
21+
import org.junit.jupiter.params.ParameterizedTest;
22+
import org.junit.jupiter.params.provider.Arguments;
23+
import org.junit.jupiter.params.provider.MethodSource;
24+
25+
@DisplayName("InferredProxyPropagator Tests")
26+
class InferredProxyPropagatorTests {
27+
private static final String PROXY_SYSTEM_KEY = "x-dd-proxy";
28+
private static final String PROXY_REQUEST_TIME_MS_KEY = "x-dd-proxy-request-time-ms";
29+
private static final String PROXY_PATH_KEY = "x-dd-proxy-path";
30+
private static final String PROXY_HTTP_METHOD_KEY = "x-dd-proxy-httpmethod";
31+
private static final String PROXY_DOMAIN_NAME_KEY = "x-dd-proxy-domain-name";
32+
private static final MapVisitor MAP_VISITOR = new MapVisitor();
33+
34+
private InferredProxyPropagator propagator;
35+
36+
@BeforeEach
37+
void setUp() {
38+
this.propagator = new InferredProxyPropagator();
39+
}
40+
41+
@Test
42+
@DisplayName("Should extract InferredProxySpan when valid headers are present")
43+
void testSuccessfulExtraction() {
44+
Map<String, String> headers = new HashMap<>();
45+
headers.put(PROXY_SYSTEM_KEY, "aws-apigateway");
46+
headers.put(PROXY_REQUEST_TIME_MS_KEY, "12345");
47+
headers.put(PROXY_PATH_KEY, "/foo");
48+
headers.put(PROXY_HTTP_METHOD_KEY, "GET");
49+
headers.put(PROXY_DOMAIN_NAME_KEY, "api.example.com");
50+
51+
Context context = this.propagator.extract(root(), headers, MAP_VISITOR);
52+
InferredProxySpan inferredProxySpan = fromContext(context);
53+
assertNotNull(inferredProxySpan);
54+
assertTrue(inferredProxySpan.isValid());
55+
}
56+
57+
@ParameterizedTest(name = "{0}")
58+
@MethodSource("invalidOrMissingHeadersProviderForPropagator")
59+
@DisplayName("Should not create InferredProxySpan if some critical headers are missing")
60+
void testExtractionWithMissingCriticalHeaders(String description, Map<String, String> headers) {
61+
Context rootContext = root();
62+
Context extractedOuterContext = this.propagator.extract(rootContext, headers, MAP_VISITOR);
63+
InferredProxySpan inferredProxySpan = fromContext(extractedOuterContext);
64+
assertNull(inferredProxySpan, "Invalid inferred proxy span should not be extracted");
65+
}
66+
67+
static Stream<Arguments> invalidOrMissingHeadersProviderForPropagator() { // Renamed
68+
Map<String, String> missingSystem = new HashMap<>();
69+
missingSystem.put(PROXY_REQUEST_TIME_MS_KEY, "12345");
70+
missingSystem.put(PROXY_PATH_KEY, "/foo");
71+
72+
Map<String, String> emptyValue = new HashMap<>();
73+
emptyValue.put(PROXY_SYSTEM_KEY, "");
74+
75+
Map<String, String> nullValue = new HashMap<>();
76+
nullValue.put(PROXY_SYSTEM_KEY, null);
77+
78+
Map<String, String> missingTime = new HashMap<>();
79+
missingTime.put(PROXY_SYSTEM_KEY, "aws-apigw");
80+
missingTime.put(PROXY_PATH_KEY, "/foo");
81+
82+
return Stream.of(
83+
of("PROXY_SYSTEM_KEY missing", missingSystem),
84+
of("PROXY_SYSTEM_KEY empty", emptyValue),
85+
of("PROXY_SYSTEM_KEY null", nullValue),
86+
of("PROXY_REQUEST_TIME_MS_KEY missing", missingTime));
87+
}
88+
89+
@ParametersAreNonnullByDefault
90+
private static class MapVisitor implements CarrierVisitor<Map<String, String>> {
91+
@Override
92+
public void forEachKeyValue(Map<String, String> carrier, BiConsumer<String, String> visitor) {
93+
carrier.forEach(visitor);
94+
}
95+
}
96+
}

‎internal-api/src/main/java/datadog/trace/api/Config.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -603,6 +603,7 @@
603603
import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_RESOURCE_REMOVE_TRAILING_SLASH;
604604
import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_SERVER_ERROR_STATUSES;
605605
import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_SERVER_PATH_RESOURCE_NAME_MAPPING;
606+
import static datadog.trace.api.config.TracerConfig.TRACE_INFERRED_PROXY_SERVICES_ENABLED;
606607
import static datadog.trace.api.config.TracerConfig.TRACE_KEEP_LATENCY_THRESHOLD_MS;
607608
import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_ENABLED;
608609
import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_FLUSH_INTERVAL;
@@ -836,6 +837,7 @@ public static String getHostName() {
836837
private final int traceBaggageMaxItems;
837838
private final int traceBaggageMaxBytes;
838839
private final List<String> traceBaggageTagKeys;
840+
private final boolean traceInferredProxyEnabled;
839841
private final int clockSyncPeriod;
840842
private final boolean logsInjectionEnabled;
841843

@@ -1754,6 +1756,8 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
17541756
tracePropagationExtractFirst =
17551757
configProvider.getBoolean(
17561758
TRACE_PROPAGATION_EXTRACT_FIRST, DEFAULT_TRACE_PROPAGATION_EXTRACT_FIRST);
1759+
traceInferredProxyEnabled =
1760+
configProvider.getBoolean(TRACE_INFERRED_PROXY_SERVICES_ENABLED, false);
17571761

17581762
clockSyncPeriod = configProvider.getInteger(CLOCK_SYNC_PERIOD, DEFAULT_CLOCK_SYNC_PERIOD);
17591763

@@ -3158,6 +3162,10 @@ public boolean isTracePropagationExtractFirst() {
31583162
return tracePropagationExtractFirst;
31593163
}
31603164

3165+
public boolean isInferredProxyPropagationEnabled() {
3166+
return traceInferredProxyEnabled;
3167+
}
3168+
31613169
public boolean isBaggageExtract() {
31623170
return tracePropagationStylesToExtract.contains(TracePropagationStyle.BAGGAGE);
31633171
}
@@ -5488,6 +5496,8 @@ public String toString() {
54885496
+ tracePropagationBehaviorExtract
54895497
+ ", tracePropagationExtractFirst="
54905498
+ tracePropagationExtractFirst
5499+
+ ", traceInferredProxyEnabled="
5500+
+ traceInferredProxyEnabled
54915501
+ ", clockSyncPeriod="
54925502
+ clockSyncPeriod
54935503
+ ", jmxFetchEnabled="

0 commit comments

Comments
 (0)