diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 8da026f224f..6cacd7037ff 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -956,7 +956,7 @@ public AgentSpan blackholeSpan() { @Override public AgentSpan.Context notifyExtensionStart(Object event) { - return LambdaHandler.notifyStartInvocation(event, propagationTagsFactory); + return LambdaHandler.notifyStartInvocation(this, event); } @Override diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java index 1d298169207..e11dbb8c29e 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java @@ -1,17 +1,13 @@ package datadog.trace.lambda; -import static datadog.trace.api.TracePropagationStyle.DATADOG; import static java.util.concurrent.TimeUnit.SECONDS; import com.squareup.moshi.JsonAdapter; import com.squareup.moshi.Moshi; import datadog.trace.api.DDSpanId; import datadog.trace.api.DDTags; -import datadog.trace.api.DDTraceId; -import datadog.trace.api.sampling.PrioritySampling; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import datadog.trace.core.propagation.ExtractedContext; -import datadog.trace.core.propagation.PropagationTags; +import datadog.trace.core.CoreTracer; import java.io.ByteArrayInputStream; import java.nio.charset.StandardCharsets; import java.util.Base64; @@ -44,7 +40,6 @@ public class LambdaHandler { private static final String DATADOG_INVOCATION_ERROR_MSG = "x-datadog-invocation-error-msg"; private static final String DATADOG_INVOCATION_ERROR_TYPE = "x-datadog-invocation-error-type"; private static final String DATADOG_INVOCATION_ERROR_STACK = "x-datadog-invocation-error-stack"; - private static final String DATADOG_TAGS_KEY = "x-datadog-tags"; private static final String START_INVOCATION = "/lambda/start-invocation"; private static final String END_INVOCATION = "/lambda/end-invocation"; @@ -73,8 +68,7 @@ public class LambdaHandler { private static String EXTENSION_BASE_URL = "http://127.0.0.1:8124"; - public static AgentSpan.Context notifyStartInvocation( - Object event, PropagationTags.Factory propagationTagsFactory) { + public static AgentSpan.Context notifyStartInvocation(CoreTracer tracer, Object event) { RequestBody body = RequestBody.create(jsonMediaType, writeValueAsString(event)); try (Response response = HTTP_CLIENT @@ -86,33 +80,16 @@ public static AgentSpan.Context notifyStartInvocation( .build()) .execute()) { if (response.isSuccessful()) { - final String traceID = response.headers().get(DATADOG_TRACE_ID); - final String priority = response.headers().get(DATADOG_SAMPLING_PRIORITY); - if (null != traceID && null != priority) { - int samplingPriority = PrioritySampling.UNSET; - try { - samplingPriority = Integer.parseInt(priority); - } catch (final NumberFormatException ignored) { - log.warn("could not read the sampling priority, defaulting to UNSET"); - } - log.debug( - "notifyStartInvocation success, found traceID = {} and samplingPriority = {}", - traceID, - samplingPriority); - PropagationTags propagationTags = - propagationTagsFactory.fromHeaderValue( - PropagationTags.HeaderType.DATADOG, response.headers().get(DATADOG_TAGS_KEY)); - return new ExtractedContext( - DDTraceId.from(traceID), - DDSpanId.ZERO, - samplingPriority, - null, - propagationTags, - DATADOG); - } else { - log.debug( - "could not find traceID or sampling priority in notifyStartInvocation, not injecting the context"); - } + + return tracer + .propagate() + .extract( + response.headers(), + (carrier, classifier) -> { + for (String headerName : carrier.names()) { + classifier.accept(headerName, carrier.get(headerName)); + } + }); } } catch (Throwable ignored) { log.error("could not reach the extension"); @@ -121,7 +98,6 @@ public static AgentSpan.Context notifyStartInvocation( } public static boolean notifyEndInvocation(AgentSpan span, Object result, boolean isError) { - if (null == span || null == span.getSamplingPriority()) { log.error( "could not notify the extension as the lambda span is null or no sampling priority has been found"); diff --git a/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy index bf2dc097f19..716dec06ddf 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy @@ -1,10 +1,9 @@ package datadog.trace.lambda -import datadog.trace.api.Config import datadog.trace.api.DDSpanId import datadog.trace.api.DDTags import datadog.trace.api.DDTraceId -import datadog.trace.core.propagation.PropagationTags +import datadog.trace.core.CoreTracer import datadog.trace.core.test.DDCoreSpecification import datadog.trace.core.DDSpan import com.amazonaws.services.lambda.runtime.events.SQSEvent @@ -30,8 +29,7 @@ class LambdaHandlerTest extends DDCoreSpecification { def "test start invocation success"() { given: - Config config = Mock(Config) - config.getxDatadogTagsMaxLength() >> 512 + CoreTracer ct = tracerBuilder().build() def server = httpServer { handlers { @@ -47,7 +45,7 @@ class LambdaHandlerTest extends DDCoreSpecification { LambdaHandler.setExtensionBaseUrl(server.address.toString()) when: - def objTest = LambdaHandler.notifyStartInvocation(obj, PropagationTags.factory(config)) + def objTest = LambdaHandler.notifyStartInvocation(ct, obj) then: objTest.getTraceId().toString() == traceId @@ -55,16 +53,50 @@ class LambdaHandlerTest extends DDCoreSpecification { cleanup: server.close() + ct.close() where: traceId | samplingPriority | obj "1234" | 2 | new TestObject() } + def "test start invocation with 128 bit trace ID"() { + given: + CoreTracer ct = tracerBuilder().build() + + def server = httpServer { + handlers { + post("/lambda/start-invocation") { + response + .status(200) + .addHeader("x-datadog-trace-id", "5744042798732701615") + .addHeader("x-datadog-sampling-priority", "2") + .addHeader("x-datadog-tags", "_dd.p.tid=1914fe7789eb32be") + .send() + } + } + } + LambdaHandler.setExtensionBaseUrl(server.address.toString()) + + when: + def objTest = LambdaHandler.notifyStartInvocation(ct, obj) + + then: + objTest.getTraceId().toHexString() == traceId + objTest.getSamplingPriority() == samplingPriority + + cleanup: + server.close() + ct.close() + + where: + traceId | samplingPriority | obj + "1914fe7789eb32be4fb6f07e011a6faf" | 2 | new TestObject() + } + def "test start invocation failure"() { given: - Config config = Mock(Config) - config.getxDatadogTagsMaxLength() >> 512 + CoreTracer ct = tracerBuilder().build() def server = httpServer { handlers { @@ -78,13 +110,14 @@ class LambdaHandlerTest extends DDCoreSpecification { LambdaHandler.setExtensionBaseUrl(server.address.toString()) when: - def objTest = LambdaHandler.notifyStartInvocation(obj, PropagationTags.factory(config)) + def objTest = LambdaHandler.notifyStartInvocation(ct, obj) then: objTest == expected cleanup: server.close() + ct.close() where: expected | obj