Skip to content

Commit

Permalink
Support Lambda PassThrough trace header propagation (#409)
Browse files Browse the repository at this point in the history 
* Support Lambda PassThrough trace header propagation

* Account for Sampled=0 trace header in Active mode

* Handle Sampled=0 Active Tracing case

* Ensure Parent is always copied from entity into trace header
  • Loading branch information
@majanjua-amzn
majanjua-amzn authored Aug 7, 2024
1 parent 101a1e6 commit 3c15d40
Show file tree
Hide file tree
Showing 9 changed files with 104 additions and 30 deletions.
12 changes: 11 additions & 1 deletion ...dk-apache-http/src/main/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import com.amazonaws.xray.AWSXRay;
import com.amazonaws.xray.AWSXRayRecorder;
import com.amazonaws.xray.entities.Namespace;
import com.amazonaws.xray.entities.NoOpSegment;
import com.amazonaws.xray.entities.Segment;
import com.amazonaws.xray.entities.Subsegment;
import com.amazonaws.xray.entities.TraceHeader;
Expand Down Expand Up @@ -104,7 +105,16 @@ public static void addRequestInformation(Subsegment subsegment, HttpRequest requ
Segment parentSegment = subsegment.getParentSegment();

if (subsegment.shouldPropagate()) {
request.addHeader(TraceHeader.HEADER_KEY, TraceHeader.fromEntity(subsegment).toString());
// If no-op, only propagate root trace ID to not taint sampling decision
TraceHeader t = TraceHeader.fromEntity(subsegment);
if (parentSegment instanceof NoOpSegment) {
request.addHeader(
TraceHeader.HEADER_KEY,
"Root=" + t.getRootTraceId().toString());
} else {
// This will propagate Parent and Sampled
request.addHeader(TraceHeader.HEADER_KEY, t.toString());
}
}

Map<String, Object> requestInformation = new HashMap<>();
Expand Down
5 changes: 3 additions & 2 deletions ...pache-http/src/test/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClientTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import static com.github.tomakehurst.wiremock.client.WireMock.anyUrl;
import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
import static com.github.tomakehurst.wiremock.client.WireMock.matching;
import static com.github.tomakehurst.wiremock.client.WireMock.ok;
import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
Expand Down Expand Up @@ -99,7 +100,7 @@ public void unsampledPropagation() throws Exception {

verify(getRequestedFor(urlPathEqualTo("/"))
.withHeader(TraceHeader.HEADER_KEY,
equalTo("Root=1-00000000-000000000000000000000000;Sampled=0")));
equalTo("Root=1-00000000-000000000000000000000000")));
}

@Test
Expand All @@ -114,7 +115,7 @@ public void unsampledButForcedPropagation() throws Exception {

verify(getRequestedFor(urlPathEqualTo("/"))
.withHeader(TraceHeader.HEADER_KEY,
equalTo("Root=1-67891233-abcdef012345678912345678;Sampled=0")));
matching("Root=1-67891233-abcdef012345678912345678;Parent=[a-z0-9]{16};Sampled=0")));
}

@Test
Expand Down
12 changes: 11 additions & 1 deletion ...rder-sdk-aws-sdk-v2/src/main/java/com/amazonaws/xray/interceptors/TracingInterceptor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import com.amazonaws.xray.entities.EntityDataKeys;
import com.amazonaws.xray.entities.EntityHeaderKeys;
import com.amazonaws.xray.entities.Namespace;
import com.amazonaws.xray.entities.NoOpSegment;
import com.amazonaws.xray.entities.Subsegment;
import com.amazonaws.xray.entities.TraceHeader;
import com.amazonaws.xray.handlers.config.AWSOperationHandler;
Expand Down Expand Up @@ -296,9 +297,18 @@ public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, Execu
return httpRequest;
}

// If no-op, only propagate root trace ID to not taint sampling decision
TraceHeader t = TraceHeader.fromEntity(subsegment);
if (subsegment.getParentSegment() instanceof NoOpSegment) {
return httpRequest.toBuilder().putHeader(
TraceHeader.HEADER_KEY,
"Root=" + t.getRootTraceId().toString()).build();
}

// This will propagate Parent and Sampled
return httpRequest.toBuilder().putHeader(
TraceHeader.HEADER_KEY,
TraceHeader.fromEntity(subsegment).toString()).build();
t.toString()).build();
}

@Override
Expand Down
12 changes: 11 additions & 1 deletion aws-xray-recorder-sdk-aws-sdk/src/main/java/com/amazonaws/xray/handlers/TracingHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
import com.amazonaws.xray.entities.EntityDataKeys;
import com.amazonaws.xray.entities.EntityHeaderKeys;
import com.amazonaws.xray.entities.Namespace;
import com.amazonaws.xray.entities.NoOpSegment;
import com.amazonaws.xray.entities.Subsegment;
import com.amazonaws.xray.entities.TraceHeader;
import com.amazonaws.xray.handlers.config.AWSOperationHandler;
Expand Down Expand Up @@ -192,7 +193,16 @@ public void beforeRequest(Request<?> request) {
currentSubsegment.setNamespace(Namespace.AWS.toString());

if (recorder.getCurrentSegment() != null && recorder.getCurrentSubsegment().shouldPropagate()) {
request.addHeader(TraceHeader.HEADER_KEY, TraceHeader.fromEntity(currentSubsegment).toString());
// If no-op, only propagate root trace ID to not taint sampling decision
TraceHeader t = TraceHeader.fromEntity(currentSubsegment);
if (currentSubsegment.getParentSegment() instanceof NoOpSegment) {
request.addHeader(
TraceHeader.HEADER_KEY,
"Root=" + t.getRootTraceId().toString());
} else {
// This will propagate Parent and Sampled
request.addHeader(TraceHeader.HEADER_KEY, t.toString());
}
}
}

Expand Down
5 changes: 4 additions & 1 deletion ...order-sdk-aws-sdk/src/test/java/com/amazonaws/xray/handlers/TracingHandlerLambdaTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,10 @@ public static void lambdaTestHelper(AWSXRayRecorder recorder, String name, boole
TraceHeader.SampleDecision.SAMPLED :
TraceHeader.SampleDecision.NOT_SAMPLED);
assertThat(traceHeader.getRootTraceId()).isEqualTo(subsegment.getTraceId());
assertThat(traceHeader.getParentId()).isEqualTo(subsegment.isSampled() ? serviceEntityId : null);
assertThat(traceHeader.getParentId()).isEqualTo(
subsegment.isSampled() ?
serviceEntityId :
subsegment.getParentSegment().getId());

tracingHandler.afterResponse(request, new Response(new InvokeResult(), new HttpResponse(request, null)));

Expand Down
44 changes: 26 additions & 18 deletions ...ray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,11 @@
import com.amazonaws.xray.AWSXRayRecorder;
import com.amazonaws.xray.entities.Entity;
import com.amazonaws.xray.entities.FacadeSegment;
import com.amazonaws.xray.entities.NoOpSegment;
import com.amazonaws.xray.entities.Segment;
import com.amazonaws.xray.entities.Subsegment;
import com.amazonaws.xray.entities.SubsegmentImpl;
import com.amazonaws.xray.entities.TraceHeader;
import com.amazonaws.xray.entities.TraceHeader.SampleDecision;
import com.amazonaws.xray.entities.TraceID;
import com.amazonaws.xray.exceptions.SubsegmentNotFoundException;
import com.amazonaws.xray.listeners.SegmentListener;
Expand All @@ -39,36 +39,42 @@ public class LambdaSegmentContext implements SegmentContext {
// See: https://github.com/aws/aws-xray-sdk-java/issues/251
private static final String LAMBDA_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader";

private static TraceHeader getTraceHeaderFromEnvironment() {
public static TraceHeader getTraceHeaderFromEnvironment() {
String lambdaTraceHeaderKey = System.getenv(LAMBDA_TRACE_HEADER_KEY);
return TraceHeader.fromString(lambdaTraceHeaderKey != null && lambdaTraceHeaderKey.length() > 0
? lambdaTraceHeaderKey
: System.getProperty(LAMBDA_TRACE_HEADER_PROP));
}

private static boolean isInitializing(TraceHeader traceHeader) {
return traceHeader.getRootTraceId() == null || traceHeader.getSampled() == null || traceHeader.getParentId() == null;
}

private static FacadeSegment newFacadeSegment(AWSXRayRecorder recorder, String name) {
TraceHeader traceHeader = getTraceHeaderFromEnvironment();
if (isInitializing(traceHeader)) {
logger.warn(LAMBDA_TRACE_HEADER_KEY + " is missing a trace ID, parent ID, or sampling decision. Subsegment "
+ name + " discarded.");
return new FacadeSegment(recorder, TraceID.create(recorder), "", SampleDecision.NOT_SAMPLED);
}
return new FacadeSegment(recorder, traceHeader.getRootTraceId(), traceHeader.getParentId(), traceHeader.getSampled());
}

// SuppressWarnings is needed for passing Root TraceId to noOp segment
@SuppressWarnings("nullness")
@Override
public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
if (logger.isDebugEnabled()) {
logger.debug("Beginning subsegment named: " + name);
}

TraceHeader traceHeader = LambdaSegmentContext.getTraceHeaderFromEnvironment();
logger.warn("TRACE HEADER IN CODE: " + traceHeader.toString());
Entity entity = getTraceEntity();
if (entity == null) { // First subsgment of a subsegment branch.
Segment parentSegment = newFacadeSegment(recorder, name);
if (entity == null) { // First subsegment of a subsegment branch
Segment parentSegment;
// Trace header either takes the structure `Root=...;<extra-data>` or
// `Root=...;Parent=...;Sampled=...;<extra-data>`
if (traceHeader.getRootTraceId() != null && traceHeader.getParentId() != null && traceHeader.getSampled() != null) {
logger.warn("CREATING FACADE SEGMENT");
parentSegment = new FacadeSegment(
recorder,
traceHeader.getRootTraceId(),
traceHeader.getParentId(),
traceHeader.getSampled());
} else {
if (logger.isDebugEnabled()) {
logger.debug("Creating No-Op parent segment");
}
TraceID t = traceHeader.getRootTraceId() != null ? traceHeader.getRootTraceId() : TraceID.create(recorder);
parentSegment = Segment.noOp(t, recorder);
}

boolean isRecording = parentSegment.isRecording();

Expand Down Expand Up @@ -145,6 +151,8 @@ public void endSubsegment(AWSXRayRecorder recorder) {
current.getCreator().getEmitter().sendSubsegment((Subsegment) current);
}
clearTraceEntity();
} else if (parentEntity instanceof NoOpSegment) {
clearTraceEntity();
} else {
setTraceEntity(current.getParent());
}
Expand Down
2 changes: 1 addition & 1 deletion aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSegment.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import java.util.concurrent.locks.ReentrantLock;
import org.checkerframework.checker.nullness.qual.Nullable;

class NoOpSegment implements Segment {
public class NoOpSegment implements Segment {

private final TraceID traceId;
private final AWSXRayRecorder creator;
Expand Down
9 changes: 8 additions & 1 deletion aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/TraceHeader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,9 +92,16 @@ public TraceHeader(@Nullable TraceID rootTraceId, @Nullable String parentId, Sam
}

public static TraceHeader fromEntity(Entity entity) {
String parentId = null;
if (entity instanceof Subsegment) {
Segment segment = entity.getParentSegment();
if (segment != null) {
parentId = segment.getId();
}
}
return new TraceHeader(
entity.getTraceId(),
entity.isSampled() ? entity.getId() : null,
(entity.getId() == null || entity.getId() == "") ? parentId : entity.getId(),
entity.isSampled() ? SampleDecision.SAMPLED : SampleDecision.NOT_SAMPLED);
}

Expand Down
33 changes: 29 additions & 4 deletions ...recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import com.amazonaws.xray.AWSXRayRecorderBuilder;
import com.amazonaws.xray.emitters.Emitter;
import com.amazonaws.xray.entities.FacadeSegment;
import com.amazonaws.xray.entities.NoOpSegment;
import com.amazonaws.xray.entities.Subsegment;
import com.amazonaws.xray.exceptions.SubsegmentNotFoundException;
import com.amazonaws.xray.strategy.LogErrorContextMissingStrategy;
Expand Down Expand Up @@ -49,6 +50,10 @@ class LambdaSegmentContextTest {

private static final String MALFORMED_TRACE_HEADER =
";;Root=1-57ff426a-80c11c39b0c928905eb0828d;;Parent=1234abcd1234abcd;;;Sampled=1;;;";
private static final String MALFORMED_TRACE_HEADER_2 = ";;root-missing;;Parent=1234abcd1234abcd;;;Sampled=1;;;";

private static final String ROOT_LAMBDA_PASSTHROUGH_TRACE_HEADER =
"Root=1-5759e988-bd862e3fe1be46a994272711;Lineage=10:1234abcd:3";

@BeforeEach
public void setupAWSXRay() {
Expand All @@ -63,14 +68,14 @@ public void setupAWSXRay() {
}

@Test
void testBeginSubsegmentWithNullTraceHeaderEnvironmentVariableResultsInAFacadeSegmentParent() {
testContextResultsInFacadeSegmentParent();
void testBeginSubsegmentWithNullTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
testContextResultsInNoOpSegmentParent();
}

@Test
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = "a")
void testBeginSubsegmentWithIncompleteTraceHeaderEnvironmentVariableResultsInAFacadeSegmentParent() {
testContextResultsInFacadeSegmentParent();
void testBeginSubsegmentWithIncompleteTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
testContextResultsInNoOpSegmentParent();
}

@Test
Expand All @@ -85,6 +90,18 @@ void testBeginSubsegmentWithCompleteButMalformedTraceHeaderEnvironmentVariableRe
testContextResultsInFacadeSegmentParent();
}

@Test
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = MALFORMED_TRACE_HEADER_2)
void testBeginSubsegmentWithIncompleteAndMalformedTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
testContextResultsInNoOpSegmentParent();
}

@Test
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = ROOT_LAMBDA_PASSTHROUGH_TRACE_HEADER)
void testBeginSubsegmentWithRootLambdaPassthroughTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
testContextResultsInNoOpSegmentParent();
}

@Test
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER_2)
void testNotSampledSetsParentToSubsegment() {
Expand Down Expand Up @@ -149,4 +166,12 @@ private static void testContextResultsInFacadeSegmentParent() {
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
assertThat(AWSXRay.getTraceEntity()).isNull();
}

private static void testContextResultsInNoOpSegmentParent() {
LambdaSegmentContext mockContext = new LambdaSegmentContext();
assertThat(mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test").getParent())
.isInstanceOf(NoOpSegment.class);
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
assertThat(AWSXRay.getTraceEntity()).isNull();
}
}

0 comments on commit 3c15d40

Please sign in to comment.