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

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

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

private void populateAndEndSubsegment(
    Subsegment currentSubsegment, Request<?> request, Response<?> response, AmazonServiceException ase) {
    if (null != response) {
        populateAndEndSubsegment(currentSubsegment, request, response);
        return;
    } else if (null != ase) {
        if (null != ase.getRequestId()) {
            currentSubsegment.putAws(REQUEST_ID_SUBSEGMENT_KEY, ase.getRequestId());
        }
        if (null != ase.getHttpHeaders() &&
            null != ase.getHttpHeaders().get(EntityHeaderKeys.AWS.EXTENDED_REQUEST_ID_HEADER)) {
            currentSubsegment.putAws(EntityDataKeys.AWS.EXTENDED_REQUEST_ID_KEY,
                                     ase.getHttpHeaders().get(EntityHeaderKeys.AWS.EXTENDED_REQUEST_ID_HEADER));
        }
        if (null != ase.getErrorMessage()) {
            currentSubsegment.getCause().setMessage(ase.getErrorMessage());
        }
        currentSubsegment.putAllHttp(extractHttpResponseInformation(ase));
    }

    finalizeSubsegment(request);
}
 

private void populateSubsegmentWithResponse(Subsegment subsegment, SdkHttpResponse httpResponse) {
    if (subsegment == null || httpResponse == null) {
        return;
    }

    String extendedRequestId = extractExtendedRequestIdFromHttp(httpResponse);
    if (extendedRequestId != null) {
        subsegment.putAws(EntityDataKeys.AWS.EXTENDED_REQUEST_ID_KEY, extendedRequestId);
    }
    subsegment.putAllHttp(extractHttpResponseParameters(httpResponse));
}
 

@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 populateAndEndSubsegment(Subsegment currentSubsegment, Request<?> request, Response<?> response) {
    if (null != response) {
        String requestId = null;
        if (response.getAwsResponse() instanceof AmazonWebServiceResult<?>) {
            // Not all services return responses extending AmazonWebServiceResult (e.g. S3)
            ResponseMetadata metadata = ((AmazonWebServiceResult<?>) response.getAwsResponse()).getSdkResponseMetadata();
            if (null != metadata) {
                requestId = metadata.getRequestId();
                if (null != requestId) {
                    currentSubsegment.putAws(REQUEST_ID_SUBSEGMENT_KEY, requestId);
                }
            }
        } else if (null != response.getHttpResponse()) { // S3 does not follow request id header convention
            if (null != response.getHttpResponse().getHeader(S3_REQUEST_ID_HEADER_KEY)) {
                currentSubsegment.putAws(REQUEST_ID_SUBSEGMENT_KEY,
                                         response.getHttpResponse().getHeader(S3_REQUEST_ID_HEADER_KEY));
            }
            if (null != response.getHttpResponse().getHeader(EntityHeaderKeys.AWS.EXTENDED_REQUEST_ID_HEADER)) {
                currentSubsegment.putAws(EntityDataKeys.AWS.EXTENDED_REQUEST_ID_KEY,
                                         response.getHttpResponse().getHeader(
                                             EntityHeaderKeys.AWS.EXTENDED_REQUEST_ID_HEADER));
            }
        }
        currentSubsegment.putAllAws(extractResponseParameters(request, response.getAwsResponse()));
        currentSubsegment.putAllHttp(extractHttpResponseInformation(response.getHttpResponse()));
    }

    finalizeSubsegment(request);
}