com.amazonaws.xray.entities.Subsegment Java Examples

/**
 * @return the current segment, or {@code Optional.empty()} if there is no segment
 */
public Optional<Segment> getCurrentSegmentOptional() {
    // explicitly do not throw context missing exceptions from optional-returning methods
    SegmentContext context = segmentContextResolverChain.resolve();
    if (null == context) {
        return Optional.empty();
    }
    Entity current = context.getTraceEntity();
    if (current instanceof Segment) {
        return Optional.of((Segment) current);
    } else if (current instanceof Subsegment) {
        return Optional.of(current.getParentSegment());
    } else {
        return Optional.empty();
    }
}
 

@Test
public void testStreamSomeChildrenRemovedFromParent() {
    TraceID traceId = new TraceID();
    DefaultStreamingStrategy defaultStreamingStrategy = new DefaultStreamingStrategy(1);

    Segment bigSegment = new SegmentImpl(AWSXRay.getGlobalRecorder(), "big", traceId);
    bigSegment.setStartTime(1.0);

    for (int i = 0; i < 5; i++) {
        Subsegment subsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "child" + i, bigSegment);
        subsegment.setStartTime(1.0);
        bigSegment.addSubsegment(subsegment);
        subsegment.end();
    }
    Assert.assertTrue(defaultStreamingStrategy.requiresStreaming(bigSegment));
    defaultStreamingStrategy.streamSome(bigSegment, AWSXRay.getGlobalRecorder().getEmitter());
    Assert.assertTrue(bigSegment.getTotalSize().intValue() == 0);
}
 

@Override
public CloseableHttpResponse execute(
        final HttpHost target,
        final HttpRequest request) throws IOException, ClientProtocolException {
    Subsegment subsegment = recorder.beginSubsegment(target.getHostName());
    return wrapHttpSupplier(subsegment, () -> {
        if (null != subsegment) {
            TracedHttpClient.addRequestInformation(subsegment, request, TracedHttpClient.getUrl(target, request));
        }
        CloseableHttpResponse response = wrappedClient.execute(target, request);
        if (null != subsegment) {
            TracedResponseHandler.addResponseInformation(subsegment, response);
        }
        return response;
    });
}
 

@Override
public CloseableHttpResponse execute(
        final HttpUriRequest request,
        final HttpContext context) throws IOException, ClientProtocolException {
    Subsegment subsegment = recorder.beginSubsegment(determineTarget(request).getHostName());
    return wrapHttpSupplier(subsegment, () -> {
        if (null != subsegment) {
            TracedHttpClient.addRequestInformation(subsegment, request, TracedHttpClient.getUrl(request));
        }
        CloseableHttpResponse response = wrappedClient.execute(request, context);
        if (null != subsegment) {
            TracedResponseHandler.addResponseInformation(subsegment, response);
        }
        return response;
    });
}
 

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);
}
 

@Override
public CloseableHttpResponse execute(
    HttpHost target, HttpRequest request, HttpContext context) throws IOException, ClientProtocolException {
    Subsegment subsegment = getRecorder().beginSubsegment(target.getHostName());
    try {
        if (null != subsegment) {
            TracedHttpClient.addRequestInformation(subsegment, request, TracedHttpClient.getUrl(target, request));
        }
        CloseableHttpResponse response = super.execute(target, request, context);
        if (null != subsegment) {
            TracedResponseHandler.addResponseInformation(subsegment, response);
        }
        return response;
    } catch (Exception e) {
        if (null != subsegment) {
            subsegment.addException(e);
        }
        throw e;
    } finally {
        if (null != subsegment) {
            getRecorder().endSubsegment();
        }
    }
}
 

@Test
public void testSegmentWithSubsegment() throws JSONException {
    TraceID traceId = new TraceID();

    Segment segment = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", traceId);
    Subsegment subsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test", segment);
    segment.addSubsegment(subsegment);

    segment.setStartTime(1.0);
    subsegment.setStartTime(1.0);

    subsegment.end();
    segment.end();

    String expected = expectedCompletedSegmentWithSubsegment(traceId, segment.getId(), subsegment.getId(), 1.0,
                                                             subsegment.getEndTime(), segment.getEndTime()).toString();

    JSONAssert.assertEquals(expected, segment.serialize(), JSONCompareMode.NON_EXTENSIBLE);

}
 

@Test
public void testingBasicStreamingFunctionality() {
    DefaultStreamingStrategy defaultStreamingStrategy = new DefaultStreamingStrategy(1);
    TraceID traceId = new TraceID();

    Segment segment = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", traceId);
    Subsegment subsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test", segment);
    Subsegment subsegment1 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test", segment);
    segment.addSubsegment(subsegment);
    segment.addSubsegment(subsegment1);

    segment.setStartTime(1.0);
    subsegment.setStartTime(1.0);
    subsegment1.setStartTime(1.0);

    subsegment.end();

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

@Override
public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) {
    AWSXRayRecorder recorder = getRecorder();
    Entity origin = recorder.getTraceEntity();

    Subsegment subsegment = recorder.beginSubsegment(executionAttributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME));
    if (subsegment == null) {
        return;
    }
    subsegment.setNamespace(Namespace.AWS.toString());
    subsegment.putAws(EntityDataKeys.AWS.OPERATION_KEY,
                      executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME));
    Region region = executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION);
    if (region != null) {
        subsegment.putAws(EntityDataKeys.AWS.REGION_KEY, region.id());
    }
    subsegment.putAllAws(extractRequestParameters(context, executionAttributes));
    if (accountId != null) {
        subsegment.putAws(EntityDataKeys.AWS.ACCOUNT_ID_SUBSEGMENT_KEY, accountId);
    }

    recorder.setTraceEntity(origin);
    // store the subsegment in the AWS SDK's executionAttributes so it can be accessed across threads
    executionAttributes.putAttribute(entityKey, subsegment);
}
 

@Test
public void testSubsegmentFunctionExceptionWhenMissingContextIsLogged() {
    // given
    RuntimeException expectedException = new RuntimeException("To be thrown by function");
    Function<Subsegment, Void> function = (subsegment) -> {
        throw expectedException;
    };
    AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard()
                                                     .withContextMissingStrategy(new LogErrorContextMissingStrategy())
                                                     .build();

    // when
    try {
        recorder.createSubsegment("test", function);
        Assert.fail("An exception should have been thrown");
    } catch (Exception e) {
        Assert.assertEquals("Function exception was not propagated", expectedException, e);
    }
}
 

@Test
public void testManuallySetEntityEndTime() {
    Segment segment = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", new TraceID());
    Subsegment subsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test", segment);
    segment.addSubsegment(subsegment);

    double endTime = 20.0d;

    segment.setStartTime(1.0);
    subsegment.setStartTime(1.0);
    segment.setEndTime(endTime);
    subsegment.setEndTime(endTime);

    subsegment.end();
    segment.end();

    Assert.assertEquals(endTime, segment.getEndTime(), 0);
    Assert.assertEquals(endTime, subsegment.getEndTime(), 0);
}
 

@Override
public void endSubsegment(AWSXRayRecorder recorder) {
    Entity current = map.get(Thread.currentThread().getId());
    if (current instanceof Subsegment) {
        Subsegment currentSubsegment = (Subsegment) current;
        if (currentSubsegment.end()) {
            recorder.sendSegment(currentSubsegment.getParentSegment());
        } else {
            if (recorder.getStreamingStrategy().requiresStreaming(currentSubsegment.getParentSegment())) {
                recorder.getStreamingStrategy().streamSome(currentSubsegment.getParentSegment(), recorder.getEmitter());
            }
            map.put(Thread.currentThread().getId(), current.getParent());
        }
    } else {
        recorder.getContextMissingStrategy().contextMissing("Failed to end subsegment: subsegment cannot be found.",
                                                            SubsegmentNotFoundException.class);
    }
}
 

private void populateRequestId(
        Subsegment subsegment,
        Optional<SdkResponse> response,
        Optional<SdkHttpResponse> httpResponse,
        Throwable exception
) {
    String requestId = null;

    if (exception != null) {
        requestId = extractRequestIdFromThrowable(exception);
    }
    if (requestId == null || requestId.equals(UNKNOWN_REQUEST_ID)) {
        requestId = extractRequestIdFromResponse(response);
    }
    if (requestId == null || requestId.equals(UNKNOWN_REQUEST_ID)) {
        requestId = extractRequestIdFromHttp(httpResponse);
    }
    if (requestId != null && !requestId.equals(UNKNOWN_REQUEST_ID)) {
        subsegment.putAws(EntityDataKeys.AWS.REQUEST_ID_KEY, requestId);
    }
}
 

/**
 * @throws SegmentNotFoundException
 *             if {@code contextMissingStrategy} throws exceptions and the segment context cannot be found
 * @throws SubsegmentNotFoundException
 *             if {@code contextMissingStrategy} throws exceptions and the current segment has no subsegments in progress
 * @return the current subsegment, or {@code null} if {@code contextMissingStrategy} suppresses exceptions and the segment
 * context cannot be found or the segment has no subsegments in progress
 */
@Nullable
public Subsegment getCurrentSubsegment() {
    SegmentContext context = getSegmentContext();
    if (context == null) {
        return null;
    }
    Entity current = context.getTraceEntity();
    if (current == null) {
        contextMissingStrategy.contextMissing("No segment in progress.", SegmentNotFoundException.class);
    } else if (current instanceof Subsegment) {
        return (Subsegment) current;
    } else {
        contextMissingStrategy.contextMissing("No subsegment in progress.", SubsegmentNotFoundException.class);
    }
    return null;
}
 

@Override
@Nullable
public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
    Entity current = getTraceEntity();
    if (current == null) {
        recorder.getContextMissingStrategy().contextMissing("Failed to begin subsegment named '" + name
                                                            + "': segment cannot be found.", SegmentNotFoundException.class);
        return null;
    }
    if (logger.isDebugEnabled()) {
        logger.debug("Beginning subsegment named: " + name);
    }
    Segment parentSegment = current.getParentSegment();
    Subsegment subsegment = new SubsegmentImpl(recorder, name, parentSegment);
    subsegment.setParent(current);
    current.addSubsegment(subsegment);
    setTraceEntity(subsegment);

    List<SegmentListener> segmentListeners = recorder.getSegmentListeners();
    segmentListeners.stream()
            .filter(Objects::nonNull)
            .forEach(listener -> listener.onBeginSubsegment(subsegment));

    return subsegment;
}
 

@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);
}
 

@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 testSubsegmentConsumerExceptionWhenMissingContextIsLogged() {
    // given
    RuntimeException expectedException = new RuntimeException("To be thrown by consumer");
    Consumer<Subsegment> consumer = (subsegment) -> {
        throw expectedException;
    };
    AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard()
                                                     .withContextMissingStrategy(new LogErrorContextMissingStrategy())
                                                     .build();

    // when
    try {
        recorder.createSubsegment("test", consumer);
        Assert.fail("An exception should have been thrown");
    } catch (Exception e) {
        Assert.assertEquals("Consumer exception was not propagated", expectedException, e);
    }
}
 

/**
 * Begins a subsegment and passes it to the supplied consumer, and ends the subsegment before returning the consumer's result.
 * Intercepts exceptions, adds them to the subsegment, and re-throws them.
 *
 * @param name
 *            the name to use for the created subsegment
 * @param consumer
 *            the function to invoke
 */
public void createSubsegment(String name, Consumer<@Nullable Subsegment> consumer) {
    Subsegment subsegment = beginSubsegment(name);
    try {
        consumer.accept(subsegment);
    } catch (Exception e) {
        if (subsegment != null) {
            subsegment.addException(e);
        }
        throw e;
    } finally {
        endSubsegment();
    }
}
 

private <R> R wrapHttpSupplier(Subsegment subsegment, HttpSupplier<R> supplier) throws IOException, ClientProtocolException {
    try {
        return supplier.get();
    } catch (Exception e) {
        if (null != subsegment) {
            subsegment.addException(e);
        }
        throw e;
    } finally {
        if (null != subsegment) {
            recorder.endSubsegment();
        }
    }
}
 

/**
 * Begins a subsegment, passes it to the provided supplier, and ends the subsegment before returning the supplier's result.
 * Intercepts exceptions, adds them to the subsegment, and re-throws them.
 *
 * @param <R>
 *            the type of the value returned by {@code function}
 * @param name
 *            the name to use for the created subsegment
 * @param supplier
 *            the supplier to invoke
 * @return the value returned by the provided supplier
 */
@Nullable
public <R> R createSubsegment(String name, Supplier<R> supplier) {
    Subsegment subsegment = beginSubsegment(name);
    try {
        return supplier.get();
    } catch (Exception e) {
        if (subsegment != null) {
            subsegment.addException(e);
        }
        throw e;
    } finally {
        endSubsegment();
    }
}
 

@Test
public void testNestedSubsegmentInjection() {
    SLF4JSegmentListener listener = (SLF4JSegmentListener) AWSXRay.getGlobalRecorder().getSegmentListeners().get(0);
    listener.setPrefix("");
    Segment seg = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", traceID);
    listener.onSetEntity(null, seg);
    Subsegment sub1 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test1", seg);
    listener.onSetEntity(seg, sub1);
    Subsegment sub2 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test2", seg);
    listener.onSetEntity(sub1, sub2);

    Assert.assertEquals(traceID.toString() + "@" + sub2.getId(), MDC.get(TRACE_ID_KEY));
}
 

@Test
public void testInjectThreadLocalInjectsCurrentSubsegment() throws Exception {
    AWSXRay.beginSegment("test");
    Subsegment subsegment = AWSXRay.beginSubsegment("test");

    threadExecutor.submit(() -> {
        AWSXRay.injectThreadLocal(subsegment);
        Assert.assertEquals(subsegment, AWSXRay.getThreadLocal());
    }).get();

    AWSXRay.endSubsegment();
    AWSXRay.endSegment();
}
 

@Benchmark
@BenchmarkMode(Mode.All)
@Fork(value=1)
@Warmup(iterations = 20)
@Measurement(iterations = 20)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
public Subsegment getSubsegmentBenchmark(PopulatedRecorderState state) {
    return state.recorder.getCurrentSubsegment();
}
 

@Benchmark
@BenchmarkMode(Mode.All)
@Fork(value=1)
@Warmup(iterations = 20)
@Measurement(iterations = 20)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
public Subsegment beginSubsegmentDummyParentBenchmark(DummyPopulatedRecorderState state) {
    return state.recorder.beginSubsegment(SUBSEGMENT_NAME);
}
 

@Benchmark
@BenchmarkMode(Mode.All)
@Fork(value=1)
@Warmup(iterations = 20)
@Measurement(iterations = 20)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
public Subsegment beginSubsegmentBenchmark(SegmentNoChildRecorderState state) {
    return state.recorder.beginSubsegment(SUBSEGMENT_NAME);
}
 

@Test
public void testSubsegmentInjection() {
    Log4JSegmentListener listener = (Log4JSegmentListener) AWSXRay.getGlobalRecorder().getSegmentListeners().get(0);
    listener.setPrefix("");
    Segment seg = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", TRACE_ID);
    listener.onSetEntity(null, seg);
    Subsegment sub = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test", seg);
    listener.onSetEntity(seg, sub);

    Assert.assertEquals(TRACE_ID.toString() + "@" + sub.getId(), ThreadContext.get(TRACE_ID_KEY));
}
 

private void traceS3() {
    // Add subsegment to current request to track call to S3
    Subsegment subsegment = AWSXRay.beginSubsegment("## Getting object metadata");
    try {
        // Gets metadata about this sample app object in S3
        s3Client.getObjectMetadata(BUCKET_NAME, OBJECT_KEY);
    } catch (Exception ex) {
        subsegment.addException(ex);
    } finally {
        AWSXRay.endSubsegment();
    }
}
 

@Test
public void testNestedSubsegmentInjection() {
    Log4JSegmentListener listener = (Log4JSegmentListener) AWSXRay.getGlobalRecorder().getSegmentListeners().get(0);
    listener.setPrefix("");
    Segment seg = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", TRACE_ID);
    listener.onSetEntity(null, seg);
    Subsegment sub1 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test1", seg);
    listener.onSetEntity(seg, sub1);
    Subsegment sub2 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test2", seg);
    listener.onSetEntity(sub1, sub2);

    Assert.assertEquals(TRACE_ID.toString() + "@" + sub2.getId(), ThreadContext.get(TRACE_ID_KEY));
}
 

@Test
public void testGetThreadLocalReturnsCurrentSubsegment() {
    AWSXRay.beginSegment("test");
    Subsegment subsegment = AWSXRay.beginSubsegment("test");
    Assert.assertEquals(subsegment, AWSXRay.getTraceEntity());
    AWSXRay.endSubsegment();
    AWSXRay.endSegment();
}