diff --git a/aws-xray-recorder-sdk-apache-http/src/main/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClient.java b/aws-xray-recorder-sdk-apache-http/src/main/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClient.java index dde253d5..71bc3dce 100644 --- a/aws-xray-recorder-sdk-apache-http/src/main/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClient.java +++ b/aws-xray-recorder-sdk-apache-http/src/main/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClient.java @@ -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; @@ -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 requestInformation = new HashMap<>(); diff --git a/aws-xray-recorder-sdk-apache-http/src/test/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClientTest.java b/aws-xray-recorder-sdk-apache-http/src/test/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClientTest.java index 1fab361a..48c0eaea 100644 --- a/aws-xray-recorder-sdk-apache-http/src/test/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClientTest.java +++ b/aws-xray-recorder-sdk-apache-http/src/test/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClientTest.java @@ -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; @@ -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 @@ -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 diff --git a/aws-xray-recorder-sdk-aws-sdk-v2/src/main/java/com/amazonaws/xray/interceptors/TracingInterceptor.java b/aws-xray-recorder-sdk-aws-sdk-v2/src/main/java/com/amazonaws/xray/interceptors/TracingInterceptor.java index 1c0ed819..027d81d5 100644 --- a/aws-xray-recorder-sdk-aws-sdk-v2/src/main/java/com/amazonaws/xray/interceptors/TracingInterceptor.java +++ b/aws-xray-recorder-sdk-aws-sdk-v2/src/main/java/com/amazonaws/xray/interceptors/TracingInterceptor.java @@ -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; @@ -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 diff --git a/aws-xray-recorder-sdk-aws-sdk/src/main/java/com/amazonaws/xray/handlers/TracingHandler.java b/aws-xray-recorder-sdk-aws-sdk/src/main/java/com/amazonaws/xray/handlers/TracingHandler.java index b9db7a17..166109b9 100644 --- a/aws-xray-recorder-sdk-aws-sdk/src/main/java/com/amazonaws/xray/handlers/TracingHandler.java +++ b/aws-xray-recorder-sdk-aws-sdk/src/main/java/com/amazonaws/xray/handlers/TracingHandler.java @@ -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; @@ -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()); + } } } diff --git a/aws-xray-recorder-sdk-aws-sdk/src/test/java/com/amazonaws/xray/handlers/TracingHandlerLambdaTest.java b/aws-xray-recorder-sdk-aws-sdk/src/test/java/com/amazonaws/xray/handlers/TracingHandlerLambdaTest.java index 1ca05b71..e1382b43 100644 --- a/aws-xray-recorder-sdk-aws-sdk/src/test/java/com/amazonaws/xray/handlers/TracingHandlerLambdaTest.java +++ b/aws-xray-recorder-sdk-aws-sdk/src/test/java/com/amazonaws/xray/handlers/TracingHandlerLambdaTest.java @@ -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))); diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java index 25c99506..6198eebe 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java +++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java @@ -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; @@ -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=...;` or + // `Root=...;Parent=...;Sampled=...;` + 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(); @@ -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()); } diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSegment.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSegment.java index e406e9ac..d470d2df 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSegment.java +++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSegment.java @@ -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; diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/TraceHeader.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/TraceHeader.java index 81700e59..427ca375 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/TraceHeader.java +++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/TraceHeader.java @@ -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); } diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java index 5f62747e..2f7a728b 100644 --- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java +++ b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java @@ -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; @@ -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() { @@ -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 @@ -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() { @@ -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(); + } }