com.amazonaws.xray.entities.TraceHeader Java Examples

@Test
@DisplayName("succeed on SpanContext extraction (TraceID)")
void contextExtractTraceId() {
    final TextMap textMap = new TextMapAdapter(Collections.singletonMap(
        TraceHeader.HEADER_KEY,
        traceHeader.toString()
    ));

    final SpanContext context = tracer.extract(Format.Builtin.TEXT_MAP, textMap);
    assertTrue(context.baggageItems().iterator().hasNext());

    final TraceHeader extractedTraceHeader = TraceHeader.fromString(context.baggageItems().iterator().next().getValue());
    assertEquals(traceHeader.getRootTraceId(), extractedTraceHeader.getRootTraceId());
    assertEquals(traceHeader.getParentId(), extractedTraceHeader.getParentId());
    assertEquals(traceHeader.getSampled(), extractedTraceHeader.getSampled());
}
 

@Test
public void testLeakedSubsegmentsAreCleanedBetweenInvocations() {

    LambdaSegmentContext lsc = new LambdaSegmentContext();

    PowerMockito.stub(PowerMockito.method(LambdaSegmentContext.class, "getTraceHeaderFromEnvironment"))
                .toReturn(TraceHeader.fromString(TRACE_HEADER));
    Subsegment firstInvocation = lsc.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
    Assert.assertNotNull(AWSXRay.getTraceEntity());

    PowerMockito.stub(PowerMockito.method(LambdaSegmentContext.class, "getTraceHeaderFromEnvironment"))
                .toReturn(TraceHeader.fromString(TRACE_HEADER_2));
    Subsegment secondInvocation = lsc.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
    Assert.assertNotNull(AWSXRay.getTraceEntity());

    Assert.assertTrue(FacadeSegment.class.isInstance(firstInvocation.getParent()));
    Assert.assertTrue(FacadeSegment.class.isInstance(secondInvocation.getParent()));
}
 

@Test
public void testSubsegmentWithChildEmittedTogetherInLambdaContext() {
    TraceHeader header = TraceHeader.fromString(TRACE_HEADER);

    PowerMockito.stub(PowerMockito.method(LambdaSegmentContext.class, "getTraceHeaderFromEnvironment")).toReturn(header);
    PowerMockito.stub(PowerMockito.method(LambdaSegmentContextResolver.class, "getLambdaTaskRoot")).toReturn("/var/task");

    Emitter mockEmitter = Mockito.mock(Emitter.class);
    AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard().withEmitter(mockEmitter).build();

    recorder.createSubsegment("testTogether", () -> {
        recorder.createSubsegment("testTogether2", () -> {
        });
    });

    ArgumentCaptor<Subsegment> emittedSubsegment = ArgumentCaptor.forClass(Subsegment.class);
    Mockito.verify(mockEmitter, Mockito.times(1)).sendSubsegment(emittedSubsegment.capture());

    Subsegment captured = emittedSubsegment.getValue();

    Assert.assertEquals(1, captured.getSubsegments().size());
}
 

@Test
public void testSubsegmentEmittedInLambdaContext() throws JSONException {
    TraceHeader header = TraceHeader.fromString(TRACE_HEADER);

    PowerMockito.stub(PowerMockito.method(LambdaSegmentContext.class, "getTraceHeaderFromEnvironment")).toReturn(header);
    PowerMockito.stub(PowerMockito.method(LambdaSegmentContextResolver.class, "getLambdaTaskRoot")).toReturn("/var/task");

    Emitter mockEmitter = Mockito.mock(Emitter.class);
    AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard().withEmitter(mockEmitter).build();
    recorder.createSubsegment("test", () -> {
    });

    ArgumentCaptor<Subsegment> emittedSubsegment = ArgumentCaptor.forClass(Subsegment.class);
    Mockito.verify(mockEmitter, Mockito.times(1)).sendSubsegment(emittedSubsegment.capture());

    Subsegment captured = emittedSubsegment.getValue();

    JSONAssert.assertEquals(expectedLambdaSubsegment(
        header.getRootTraceId(), header.getParentId(), captured.getId(), captured.getStartTime(),
        captured.getEndTime()).toString(), captured.streamSerialize(), JSONCompareMode.NON_EXTENSIBLE);
}
 

@Override
public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
    SdkHttpRequest httpRequest = context.httpRequest();

    Subsegment subsegment = executionAttributes.getAttribute(entityKey);
    if (subsegment == null) {
        return httpRequest;
    }

    boolean isSampled = subsegment.getParentSegment().isSampled();
    TraceHeader header = new TraceHeader(
            subsegment.getParentSegment().getTraceId(),
            isSampled ? subsegment.getId() : null,
            isSampled ? TraceHeader.SampleDecision.SAMPLED : TraceHeader.SampleDecision.NOT_SAMPLED
    );

    return httpRequest.toBuilder().appendHeader(TraceHeader.HEADER_KEY, header.toString()).build();
}
 

public static void addRequestInformation(Subsegment subsegment, HttpRequest request, String url) {
    subsegment.setNamespace(Namespace.REMOTE.toString());
    Segment parentSegment = subsegment.getParentSegment();

    TraceHeader header = new TraceHeader(parentSegment.getTraceId(),
                            parentSegment.isSampled() ? subsegment.getId() : null,
                            parentSegment.isSampled() ? SampleDecision.SAMPLED : SampleDecision.NOT_SAMPLED);
    request.addHeader(TraceHeader.HEADER_KEY, header.toString());

    Map<String, Object> requestInformation = new HashMap<>();

    requestInformation.put("url", url);
    requestInformation.put("method", request.getRequestLine().getMethod());

    subsegment.putHttp("request", requestInformation);
}
 

/**
 * In systems where the surrounding code is using X-Ray directly, but
 * not the OpenTracing API, we should detect if a trace is already in
 * progress. For example, in AWS Lambda functions, the lambda server
 * creates a top-level trace segment for the whole function call
 *
 * @see <a href="https://docs.aws.amazon.com/xray/latest/devguide/xray-services-lambda.html">https://docs.aws.amazon.com/xray/latest/devguide/xray-services-lambda.html</a>
 */
@Test
@DisplayName("detect a pre-existing X-Ray trace")
void detectPreExisting() {
    final Segment parentEntity = awsxRayRecorder.beginSegment("pre-existing-trace");

    final Scope childScope = tracer
            .buildSpan("child-of-pre-existing-trace")
            .startActive(true);

    final Entity childEntity = ((AWSXRayScope) childScope).span().getEntity();

    assertFalse(parentEntity.getSubsegments().isEmpty());
    assertEquals(parentEntity, childEntity.getParent());

    // Check that trace header is correctly set in the child
    final String childTraceHeader = childScope.span().getBaggageItem(TraceHeader.HEADER_KEY);
    assertNotNull(childTraceHeader);
    assertTrue(childTraceHeader.contains(parentEntity.getTraceId().toString()));
    assertTrue(childTraceHeader.contains(parentEntity.getId()));

    childScope.close();
}
 

@Test
@DisplayName("succeed on SpanContext injection (TraceID)")
void contextInjectTraceId() {
    final TextMap textMap = new TextMapAdapter(new HashMap<>());
    final SpanContext context = new AWSXRaySpanContext(Collections.singletonMap(
        TraceHeader.HEADER_KEY,
        traceHeader.toString()
    ));

    tracer.inject(context, Format.Builtin.TEXT_MAP, textMap);
    assertTrue(textMap.iterator().hasNext());

    final TraceHeader extractedTraceHeader = TraceHeader.fromString(textMap.iterator().next().getValue());
    assertEquals(traceHeader.getRootTraceId(), extractedTraceHeader.getRootTraceId());
    assertEquals(traceHeader.getParentId(), extractedTraceHeader.getParentId());
    assertEquals(traceHeader.getSampled(), extractedTraceHeader.getSampled());
}
 

@Test
public void testLongHeaderFromStringWithSpaces() {
    TraceHeader header = TraceHeader.fromString("Sampled=?; Root=" + TRACE_ID + "; Parent=foo; Self=2; Foo=bar");
    Assert.assertEquals(SampleDecision.REQUESTED, header.getSampled());
    Assert.assertEquals(TraceID.fromString(TRACE_ID), header.getRootTraceId());
    Assert.assertEquals("foo", header.getParentId());
    Assert.assertEquals(1, header.getAdditionalParams().size());
    Assert.assertEquals("bar", header.getAdditionalParams().get("Foo"));
}
 

@Override
public void beforeRequest(Request<?> request) {
    String serviceName = extractServiceName(request);
    String operationName = extractOperationName(request);

    if (S3_SERVICE_NAME.equals(serviceName) && S3_PRESIGN_REQUEST.equals(operationName)) {
        return;
    }

    if (XRAY_SERVICE_NAME.equals(serviceName) && (XRAY_SAMPLING_RULE_REQUEST.equals(operationName)
            || XRAY_SAMPLING_TARGET_REQUEST.equals(operationName))) {
        return;
    }

    if (isSubsegmentDuplicate(recorder.getCurrentSubsegmentOptional(), request)) {
        return;
    }
    Entity entityContext = request.getHandlerContext(ENTITY_KEY);
    if (null != entityContext) {
        recorder.setTraceEntity(entityContext);
    }
    Subsegment currentSubsegment = recorder.beginSubsegment(serviceName);
    if (null == currentSubsegment) {
        return;
    }
    currentSubsegment.putAllAws(extractRequestParameters(request));
    currentSubsegment.putAws(EntityDataKeys.AWS.OPERATION_KEY, operationName);
    if (null != accountId) {
        currentSubsegment.putAws(EntityDataKeys.AWS.ACCOUNT_ID_SUBSEGMENT_KEY, accountId);
    }
    currentSubsegment.setNamespace(Namespace.AWS.toString());

    if (null != recorder.getCurrentSegment()) {
        TraceHeader header =
            new TraceHeader(recorder.getCurrentSegment().getTraceId(),
                            recorder.getCurrentSegment().isSampled() ? currentSubsegment.getId() : null,
                            recorder.getCurrentSegment().isSampled() ? SampleDecision.SAMPLED : SampleDecision.NOT_SAMPLED);
        request.addHeader(TraceHeader.HEADER_KEY, header.toString());
    }
}
 

private void testMockContext(TraceHeader xAmznTraceId, Class<?> instanceOfClass) {
    LambdaSegmentContext mockContext = mockContext(xAmznTraceId);
    Assert.assertTrue(
        instanceOfClass.isInstance(mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test").getParent()));
    mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
    Assert.assertNull(AWSXRay.getTraceEntity());
}
 

@Test
public void testDefaultStreamingStrategyForLambdaTraceContext() {
    DefaultStreamingStrategy defaultStreamingStrategy = new DefaultStreamingStrategy(1);

    //if FacadeSegment size is larger than maxSegmentSize and only the first subsegment is completed, first subsegment will be
    //streamed out
    FacadeSegment facadeSegmentOne = new FacadeSegment(AWSXRay.getGlobalRecorder(), new TraceID(), "",
                                                       TraceHeader.SampleDecision.SAMPLED);
    Subsegment firstSubsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "FirstSubsegment", facadeSegmentOne);
    Subsegment secondSubsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "SecondSubsegment", facadeSegmentOne);
    facadeSegmentOne.addSubsegment(firstSubsegment);
    facadeSegmentOne.addSubsegment(secondSubsegment);

    firstSubsegment.end();

    Assert.assertTrue(facadeSegmentOne.getTotalSize().intValue() == 2);
    defaultStreamingStrategy.streamSome(facadeSegmentOne, AWSXRay.getGlobalRecorder().getEmitter());
    Assert.assertTrue(facadeSegmentOne.getTotalSize().intValue() == 1);

    Subsegment tempOne = facadeSegmentOne.getSubsegments().get(0);
    Assert.assertEquals("SecondSubsegment", tempOne.getName());

    //if FarcadeSegment size is larger than maxSegmentSize and only the second subsegment is completed, second subsegment will
    //be streamed out
    FacadeSegment facadeSegmentTwo = new FacadeSegment(AWSXRay.getGlobalRecorder(), new TraceID(), "",
                                                       TraceHeader.SampleDecision.SAMPLED);
    Subsegment thirdSubsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "ThirdSubsegment", facadeSegmentTwo);
    Subsegment fourthSubsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "FourthSubsegment", facadeSegmentTwo);
    facadeSegmentTwo.addSubsegment(thirdSubsegment);
    facadeSegmentTwo.addSubsegment(fourthSubsegment);

    fourthSubsegment.end();

    Assert.assertTrue(facadeSegmentTwo.getTotalSize().intValue() == 2);
    defaultStreamingStrategy.streamSome(facadeSegmentTwo, AWSXRay.getGlobalRecorder().getEmitter());
    Assert.assertTrue(facadeSegmentTwo.getTotalSize().intValue() == 1);

    Subsegment tempTwo = facadeSegmentTwo.getSubsegments().get(0);
    Assert.assertEquals("ThirdSubsegment", tempTwo.getName());
}
 

@Test
public void testLongHeaderToString() {
    TraceHeader header = new TraceHeader();
    header.setSampled(SampleDecision.SAMPLED);
    header.setRootTraceId(TraceID.fromString(TRACE_ID));
    header.setParentId("foo");
    header.getAdditionalParams().put("Foo", "bar");
    Assert.assertEquals("Root=" + TRACE_ID + ";Parent=foo;Sampled=1;Foo=bar", header.toString());
}
 

@Test
public void testSampledEqualsOneAndParentToString() {
    TraceHeader header = new TraceHeader();
    header.setSampled(SampleDecision.SAMPLED);
    header.setParentId("foo");
    Assert.assertEquals("Parent=foo;Sampled=1", header.toString());
}
 

@Test
public void testSampledEqualsOneWithSamplingRequestedToString() {
    TraceHeader header = new TraceHeader();
    header.setSampled(SampleDecision.REQUESTED);
    header.setSampled(SampleDecision.SAMPLED);
    Assert.assertEquals("Sampled=1", header.toString());
}
 

@Test
@DisplayName("ignore implicit active span on ignoreActiveSpan")
void ignoreImplicitParentSpan() {
    final Scope parentScope = tracer
            .buildSpan("parent-span")
            .startActive(true);

    final Scope childScope = tracer
            .buildSpan("child-span")
            .ignoreActiveSpan()
            .startActive(true);

    final Entity parentEntity = ((AWSXRayScope) parentScope).span().getEntity();
    final Entity childEntity = ((AWSXRayScope) childScope).span().getEntity();

    assertTrue(parentEntity.getSubsegments().isEmpty());
    assertNull(childEntity.getParent());
    assertNotEquals(parentEntity.getParentSegment().getTraceId(), childEntity.getParentSegment().getTraceId());

    // Check that trace header is correctly set in the child
    final String childTraceHeader = childScope.span().getBaggageItem(TraceHeader.HEADER_KEY);
    assertNotNull(childTraceHeader);
    assertTrue(childTraceHeader.contains(childEntity.getParentSegment().getTraceId().toString()));
    assertFalse(childTraceHeader.contains(parentEntity.getParentSegment().getTraceId().toString()));
    assertFalse(childTraceHeader.contains(parentEntity.getId()));

    childScope.close();
    parentScope.close();
}
 

@Test
public void testLongHeaderFromString() {
    TraceHeader header = TraceHeader.fromString("Sampled=?;Root=" + TRACE_ID + ";Parent=foo;Self=2;Foo=bar");
    Assert.assertEquals(SampleDecision.REQUESTED, header.getSampled());
    Assert.assertEquals(TraceID.fromString(TRACE_ID), header.getRootTraceId());
    Assert.assertEquals("foo", header.getParentId());
    Assert.assertEquals(1, header.getAdditionalParams().size());
    Assert.assertEquals("bar", header.getAdditionalParams().get("Foo"));
}
 

@Test
public void testSampledEqualsOneFromString() {
    TraceHeader header = TraceHeader.fromString("Sampled=1");
    Assert.assertEquals(SampleDecision.SAMPLED, header.getSampled());
    Assert.assertNull(header.getRootTraceId());
    Assert.assertNull(header.getParentId());
    Assert.assertTrue(header.getAdditionalParams().isEmpty());
}
 

@Test
@DisplayName("set trace header in baggage")
void setTraceHeaderInBaggage() {
    final Scope activeScope = tracer
            .buildSpan("test-trace-header")
            .startActive(true);

    final String activeTraceHeader = activeScope.span().getBaggageItem(TraceHeader.HEADER_KEY);
    assertNotNull(activeTraceHeader);
    assertTrue(activeTraceHeader.contains(((AWSXRaySpan) activeScope.span()).getEntity().getTraceId().toString()));

    activeScope.close();
}
 

@Test
@DisplayName("set explicit span as parent from remote server")
void setExplicitParentSpanFromRemote() {

    // SpanContext can be passed to remote servers using inject() and
    // extract(), so assume we read this in from e.g. HTTP headers
    final SpanContext remoteContext = new AWSXRaySpanContext(Collections.singletonMap(
            TraceHeader.HEADER_KEY,
            traceHeader.toString()
    ));

    final Scope childScope =  tracer
            .buildSpan("child-span")
            .asChildOf(remoteContext)
            .startActive(true);

    final Entity childEntity  = ((AWSXRayScope) childScope).span().getEntity();

    assertEquals(childEntity.getParentSegment().getTraceId(), traceHeader.getRootTraceId());
    assertEquals(childEntity.getParentSegment().getId(), traceHeader.getParentId());

    // Check that trace header is correctly set in the child
    final String childTraceHeader = childScope.span().getBaggageItem(TraceHeader.HEADER_KEY);
    assertNotNull(childTraceHeader);

    childScope.close();
}
 

@Test
@DisplayName("set explicit span as parent")
void setExplicitParentSpan() {

    // NB we *don't* startActive here - assume this Span
    // object came from somewhere else in the code
    final AWSXRaySpan explicitParentSpan = mockSpan("explicit-parent-span");

    // This implicit parent should be ignored by SpanBuilder
    // when we set the explicit parent
    final Scope implicitParentScope = tracer
            .buildSpan("implicit-parent-span")
            .startActive(true);

    final Scope childScope =  tracer
            .buildSpan("child-span")
            .asChildOf(explicitParentSpan)
            .startActive(true);

    final Entity explicitParentEntity = explicitParentSpan.getEntity();
    final Entity implicitParentEntity = ((AWSXRayScope) implicitParentScope).span().getEntity();
    final Entity childEntity  = ((AWSXRayScope) childScope).span().getEntity();

    assertFalse(explicitParentEntity.getSubsegments().isEmpty());
    assertTrue(implicitParentEntity.getSubsegments().isEmpty());

    assertEquals(explicitParentEntity, childEntity.getParent());
    assertNotEquals(explicitParentEntity.getId(), childEntity.getId());

    // Check that trace header is correctly set in the child
    final String childTraceHeader = childScope.span().getBaggageItem(TraceHeader.HEADER_KEY);
    assertNotNull(childTraceHeader);
    assertTrue(childTraceHeader.contains(explicitParentEntity.getTraceId().toString()));
    assertTrue(childTraceHeader.contains(explicitParentEntity.getId()));

    childScope.close();
    implicitParentScope.close();
}
 

private Optional<TraceHeader> getTraceHeader(HttpServletRequest request) {
    String traceHeaderString = request.getHeader(TraceHeader.HEADER_KEY);
    if (null != traceHeaderString) {
        return Optional.of(TraceHeader.fromString(traceHeaderString));
    }
    return Optional.empty();
}
 

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(), "", SampleDecision.NOT_SAMPLED);
    }
    return new FacadeSegment(recorder, traceHeader.getRootTraceId(), traceHeader.getParentId(), traceHeader.getSampled());
}
 

@Test
@DisplayName("set implicit active span as parent")
void setImplicitParentSpan() {
    final Scope parentScope = tracer
            .buildSpan("parent-span")
            .startActive(true);

    final Scope childScope =  tracer
            .buildSpan("child-span")
            .startActive(true);

    final Entity parentEntity = ((AWSXRayScope) parentScope).span().getEntity();
    final Entity childEntity  = ((AWSXRayScope) childScope).span().getEntity();

    assertFalse(parentEntity.getSubsegments().isEmpty());
    assertEquals(parentEntity, childEntity.getParent());
    assertEquals(parentEntity.getTraceId(), childEntity.getParent().getTraceId());

    // Check that trace header is correctly set in the child
    final String childTraceHeader = childScope.span().getBaggageItem(TraceHeader.HEADER_KEY);
    assertNotNull(childTraceHeader);
    assertTrue(childTraceHeader.contains(parentEntity.getTraceId().toString()));
    assertTrue(childTraceHeader.contains(parentEntity.getId()));

    childScope.close();
    parentScope.close();
}
 

@Test
public void testSubsegmentNotEmittedWithoutExceptionInLambdaInitContext() {
    PowerMockito.stub(PowerMockito.method(LambdaSegmentContext.class, "getTraceHeaderFromEnvironment"))
                .toReturn(TraceHeader.fromString(null));
    PowerMockito.stub(PowerMockito.method(LambdaSegmentContextResolver.class, "getLambdaTaskRoot")).toReturn("/var/task");

    Emitter mockEmitter = Mockito.mock(Emitter.class);
    AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard().withEmitter(mockEmitter).build();
    recorder.createSubsegment("test", () -> {
    });

    Mockito.verify(mockEmitter, Mockito.times(0)).sendSubsegment(Mockito.any());
}
 

@Test
public void testSubsequentSubsegmentBranchesEmittedInLambdaContext() {
    TraceHeader header = TraceHeader.fromString(TRACE_HEADER);

    PowerMockito.stub(PowerMockito.method(LambdaSegmentContext.class, "getTraceHeaderFromEnvironment")).toReturn(header);
    PowerMockito.stub(PowerMockito.method(LambdaSegmentContextResolver.class, "getLambdaTaskRoot")).toReturn("/var/task");

    Emitter mockEmitter = Mockito.mock(Emitter.class);
    AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard().withEmitter(mockEmitter).build();

    recorder.createSubsegment("testTogether", () -> {
        recorder.createSubsegment("testTogether2", () -> {
        });
    });

    recorder.createSubsegment("testTogether3", () -> {
        recorder.createSubsegment("testTogether4", () -> {
        });
    });

    ArgumentCaptor<Subsegment> emittedSubsegments = ArgumentCaptor.forClass(Subsegment.class);

    Mockito.verify(mockEmitter, Mockito.times(2)).sendSubsegment(emittedSubsegments.capture());

    List<Subsegment> captured = emittedSubsegments.getAllValues();

    captured.forEach((capturedSubsegment) -> {
        Assert.assertEquals(1, capturedSubsegment.getSubsegments().size());
    });
}
 

@Test
public void testSampledUnknownToString() {
    TraceHeader header = new TraceHeader();
    header.setSampled(SampleDecision.UNKNOWN);
    Assert.assertEquals("", header.toString());
}
 

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

private static TraceHeader getTraceHeaderFromEnvironment() {
    return TraceHeader.fromString(System.getenv(LAMBDA_TRACE_HEADER_KEY));
}
 

@Test
public void testBeginSubsegmentWithNullTraceHeaderEnvironmentVariableResultsInADummySegmentParent() {
    testMockContext(TraceHeader.fromString(null), FacadeSegment.class);
}