Java Code Examples for com.amazonaws.xray.entities.Subsegment#end()

/**
 * Ends the provided subsegment. This method doesn't touch context storage.
 *
 * @param subsegment
 *          the subsegment to close.
 */
public void endSubsegment(Subsegment subsegment) {
    if (subsegment == null) {
        logger.debug("No input subsegment to end. No-op.");
        return;
    }
    boolean rootReady = subsegment.end();
    // First handling the special case where its direct parent is a facade segment
    if (subsegment.getParent() instanceof FacadeSegment) {
        if (((FacadeSegment) subsegment.getParent()).isSampled()) {
            getEmitter().sendSubsegment(subsegment);
        }
        return;
    }
    // Otherwise we check the happy case where the entire segment is ready
    if (rootReady && !(subsegment.getParentSegment() instanceof FacadeSegment)) {
        sendSegment(subsegment.getParentSegment());
        return;
    }
    // If not we try to stream closed subsegments regardless the root segment is facade or real
    if (this.getStreamingStrategy().requiresStreaming(subsegment.getParentSegment())) {
        this.getStreamingStrategy().streamSome(subsegment.getParentSegment(), this.getEmitter());
    }
}
 

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

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

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

@Override
public void endSubsegment(AWSXRayRecorder recorder) {
    Entity current = getTraceEntity();
    if (current instanceof Subsegment) {
        if (logger.isDebugEnabled()) {
            logger.debug("Ending subsegment named: " + current.getName());
        }
        Subsegment currentSubsegment = (Subsegment) current;

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

        currentSubsegment.end();

        if (recorder.getStreamingStrategy().requiresStreaming(currentSubsegment.getParentSegment())) {
            recorder.getStreamingStrategy().streamSome(currentSubsegment.getParentSegment(), recorder.getEmitter());
        }

        segmentListeners
                .stream()
                .filter(Objects::nonNull)
                .forEach(listener -> listener.afterEndSubsegment(currentSubsegment));

        Entity parentEntity = current.getParent();
        if (parentEntity instanceof FacadeSegment) {
            if (((FacadeSegment) parentEntity).isSampled()) {
                current.getCreator().getEmitter().sendSubsegment((Subsegment) current);
            }
            clearTraceEntity();
        } else {
            setTraceEntity(current.getParent());
        }

    } else {
        throw new SubsegmentNotFoundException("Failed to end a subsegment: subsegment cannot be found.");
    }
}
 

@Override
public void endSubsegment(AWSXRayRecorder recorder) {
    Entity current = getTraceEntity();
    if (current instanceof Subsegment) {
        if (logger.isDebugEnabled()) {
            logger.debug("Ending subsegment named: " + current.getName());
        }
        Subsegment currentSubsegment = (Subsegment) current;

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

        if (currentSubsegment.end()) {
            recorder.sendSegment(currentSubsegment.getParentSegment());
        } else {
            if (recorder.getStreamingStrategy().requiresStreaming(currentSubsegment.getParentSegment())) {
                recorder.getStreamingStrategy().streamSome(currentSubsegment.getParentSegment(), recorder.getEmitter());
            }

            segmentListeners
                    .stream()
                    .filter(Objects::nonNull)
                    .forEach(listener -> listener.afterEndSubsegment(currentSubsegment));

            setTraceEntity(current.getParent());
        }
    } else {
        recorder.getContextMissingStrategy().contextMissing("Failed to end subsegment: subsegment cannot be found.",
                                                            SubsegmentNotFoundException.class);
    }
}
 

@Test
public void testBushyandSpindlySegmentTreeStreaming() {
    TraceID traceId = new TraceID();

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

    SubsegmentImpl holder = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "big_child0", bigSegment);
    holder.setStartTime(1.0);
    bigSegment.addSubsegment(holder);
    holder.end();

    SubsegmentImpl holder1 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "big_child1", bigSegment);
    holder1.setStartTime(1.0);
    bigSegment.addSubsegment(holder1);
    holder1.end();

    SubsegmentImpl holder2 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "big_child2", bigSegment);
    holder2.setStartTime(1.0);
    bigSegment.addSubsegment(holder2);
    holder2.end();

    DefaultStreamingStrategy defaultStreamingStrategy = new DefaultStreamingStrategy(1);
    Assert.assertTrue(defaultStreamingStrategy.requiresStreaming(bigSegment));
    defaultStreamingStrategy.streamSome(bigSegment, AWSXRay.getGlobalRecorder().getEmitter());
    Assert.assertTrue(bigSegment.getReferenceCount() == 0);
}
 

@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 testMultithreadedStreamSome() {
    DefaultStreamingStrategy defaultStreamingStrategy = new DefaultStreamingStrategy(1);

    Segment segment = AWSXRay.beginSegment("big");


    Subsegment subsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "subsegment1", segment);
    subsegment.setStartTime(1.0);
    segment.addSubsegment(subsegment);
    subsegment.end();

    Thread thread1 = new Thread(() -> {
        AWSXRay.setTraceEntity(segment);
        AWSXRay.beginSubsegment("thread1");
        AWSXRay.endSubsegment();
    });
    Thread thread2 = new Thread(() -> {
        AWSXRay.setTraceEntity(segment);
        AWSXRay.beginSubsegment("thread2");
        AWSXRay.endSubsegment();
    });

    thread1.start();
    thread2.start();
    for (Thread thread : new Thread[]{thread1, thread2}) {
        try {
            thread.join();
        } catch (InterruptedException e) {
            return;
        }
    }

    Assert.assertTrue(AWSXRay.getTraceEntity().getName().equals("big"));
    //asserts that all subsegments are added correctly.
    Assert.assertTrue(AWSXRay.getTraceEntity().getTotalSize().intValue() == 3);

    defaultStreamingStrategy.streamSome(segment, AWSXRay.getGlobalRecorder().getEmitter());

    Assert.assertTrue(segment.getTotalSize().intValue() == 0);
}