Java Code Examples for com.amazonaws.xray.AWSXRay#endSubsegment()

private Segment segmentInResponseToCode(int code) {
    NoOpResponseHandler responseHandler = new NoOpResponseHandler();
    TracedResponseHandler<String> tracedResponseHandler = new TracedResponseHandler<>(responseHandler);
    HttpResponse httpResponse = new BasicHttpResponse(new BasicStatusLine(HttpVersion.HTTP_1_1, code, ""));

    Segment segment = AWSXRay.beginSegment("test");
    AWSXRay.beginSubsegment("someHttpCall");

    try {
        tracedResponseHandler.handleResponse(httpResponse);
    } catch (IOException e) {
        throw new RuntimeException(e);
    }

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

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

protected Object processXRayTrace(ProceedingJoinPoint pjp) throws Throwable {
    try {
        Subsegment subsegment = AWSXRay.beginSubsegment(pjp.getSignature().getName());
        if (subsegment != null) {
            subsegment.setMetadata(generateMetadata(pjp, subsegment));
        }
        return XRayInterceptorUtils.conditionalProceed(pjp);
    } catch (Exception e) {
        AWSXRay.getCurrentSegmentOptional().ifPresent(x -> x.addException(e));
        throw e;
    } finally {
        logger.trace("Ending Subsegment");
        AWSXRay.endSubsegment();
    }
}
 

@Around("queryExecution()")
public Object traceSQL(ProceedingJoinPoint pjp) throws Throwable {
    try {
        Subsegment subsegment = AWSXRay.beginSubsegment(pjp.getSignature().getName());
        XRayInterceptorUtils.generateMetadata(pjp, subsegment);
        return XRayInterceptorUtils.conditionalProceed(pjp);
    } catch (Exception e) {
        logger.error(e.getMessage());
        AWSXRay.getCurrentSegment().addException(e);
        throw e;
    } finally {
        logger.trace("Ending Subsegment");
        AWSXRay.endSubsegment();
    }
}
 

@Override
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
    Subsegment subsegment = null;

    if (isExecution(method)) {
        // only trace on execution methods
        subsegment = createSubsegment();
    }

    logger.debug(
        String.format("Invoking statement execution with X-Ray tracing. Tracing active: %s", subsegment != null));
    try {
        // execute the query "wrapped" in a XRay Subsegment
        return method.invoke(delegate, args);
    } catch (Throwable t) {
        Throwable rootThrowable = t;
        if (t instanceof InvocationTargetException) {
            // the reflection may wrap the actual error with an InvocationTargetException.
            // we want to use the root cause to make the instrumentation seamless
            InvocationTargetException ite = (InvocationTargetException) t;
            if (ite.getTargetException() != null) {
                rootThrowable = ite.getTargetException();
            } else if (ite.getCause() != null) {
                rootThrowable = ite.getCause();
            }
        }

        if (subsegment != null) {
            subsegment.addException(rootThrowable);
        }
        throw rootThrowable;
    } finally {
        if (subsegment != null && isExecution(method)) {
            AWSXRay.endSubsegment();
        }
    }
}
 

@Test
public void testSubsegmentListeners() {
    AWSXRay.beginSegment("test");
    Subsegment sub = AWSXRay.beginSubsegment("testSub");
    String beginAnnotation = sub.getAnnotations().get("subAnnotation1").toString();
    AWSXRay.endSubsegment();
    String endAnnotation = sub.getAnnotations().get("subAnnotation2").toString();

    Assert.assertEquals("began", beginAnnotation);
    Assert.assertEquals("ended", endAnnotation);
}
 

@Override
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
    //get the name of the method for comparison
    final String name = method.getName();
    //was close invoked?
    boolean close = compare(JdbcInterceptor.CLOSE_VAL, name);
    //allow close to be called multiple times
    if (close && closed) {
        return null;
    }
    //are we calling isClosed?
    if (compare(JdbcInterceptor.ISCLOSED_VAL, name)) {
        return Boolean.valueOf(closed);
    }
    //if we are calling anything else, bail out
    if (closed) {
        throw new SQLException("Statement closed.");
    }
    //check to see if we are about to execute a query
    final boolean process = isExecute(method);
    Object result = null;
    Subsegment subsegment = null;
    if (process) {
        subsegment = AWSXRay.beginSubsegment(hostname);
    }
    try {
        if (process && null != subsegment) {
            subsegment.putAllSql(additionalParams);
            subsegment.setNamespace(Namespace.REMOTE.toString());
        }
        result = method.invoke(delegate, args); //execute the query
    } catch (Throwable t) {
        if (null != subsegment) {
            subsegment.addException(t);
        }
        if (t instanceof InvocationTargetException && t.getCause() != null) {
            throw t.getCause();
        } else {
            throw t;
        }
    } finally {
        if (process && null != subsegment) {
            AWSXRay.endSubsegment();
        }
    }

    //perform close cleanup
    if (close) {
        closed = true;
        delegate = null;
    }

    return result;
}
 

@Override
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
    //get the name of the method for comparison
    final String name = method.getName();
    //was close invoked?
    boolean close = compare(JdbcInterceptor.CLOSE_VAL, name);
    //allow close to be called multiple times
    if (close && closed) {
        return null;
    }
    //are we calling isClosed?
    if (compare(JdbcInterceptor.ISCLOSED_VAL, name)) {
        return Boolean.valueOf(closed);
    }
    //if we are calling anything else, bail out
    if (closed) {
        throw new SQLException("Statement closed.");
    }
    //check to see if we are about to execute a query
    final boolean process = isExecute(method);
    Object result = null;
    Subsegment subsegment = null;
    if (process) {
        subsegment = AWSXRay.beginSubsegment(hostname);
    }
    try {
        if (process && null != subsegment) {
            subsegment.putAllSql(additionalParams);
            subsegment.setNamespace(Namespace.REMOTE.toString());
        }
        result = method.invoke(delegate, args); //execute the query
    } catch (Throwable t) {
        if (null != subsegment) {
            subsegment.addException(t);
        }
        if (t instanceof InvocationTargetException && t.getCause() != null) {
            throw t.getCause();
        } else {
            throw t;
        }
    } finally {
        if (process && null != subsegment) {
            AWSXRay.endSubsegment();
        }
    }

    //perform close cleanup
    if (close) {
        closed = true;
        delegate = null;
    }

    return result;
}
 

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