Java Code Examples for software.amazon.awssdk.core.interceptor.ExecutionAttributes#getAttribute()

/**
 * Before an individual http request is finalized. This is only called once per operation, meaning
 * we can only have one span per operation.
 */
@Override public SdkHttpRequest modifyHttpRequest(
    Context.ModifyHttpRequest context,
    ExecutionAttributes executionAttributes
) {
  HttpClientRequest request = new HttpClientRequest(context.httpRequest());

  Span span = handler.handleSend(request);
  executionAttributes.putAttribute(SPAN, span);

  String serviceName = executionAttributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME);
  String operation = getAwsOperationNameFromRequestClass(context.request());
  // TODO: This overwrites user configuration. We don't do this in other layered tools such
  // as WebMVC. Instead, we add tags (such as we do here) and neither overwrite the name, nor
  // remoteServiceName. Users can always remap in an span handler using tags!
  span.name(operation)
      .remoteServiceName(serviceName)
      .tag("aws.service_name", serviceName)
      .tag("aws.operation", operation);

  return request.build();
}
 

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes attrs) {
  logRetryAttempt(attrs);

  String serviceName = attrs.getAttribute(SdkExecutionAttribute.SERVICE_NAME);
  String opName = attrs.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
  String endpoint = serviceName + "." + opName;

  SdkHttpRequest request = context.httpRequest();

  IpcLogEntry logEntry = logger.createClientEntry()
      .withOwner("aws-sdk-java-v2")
      .withProtocol(IpcProtocol.http_1)
      .withHttpMethod(request.method().name())
      .withUri(request.getUri())
      .withEndpoint(endpoint)
      .withAttempt(extractAttempt(request))
      .withAttemptFinal(false); // Don't know if it is the final attempt

  request.headers().forEach((k, vs) -> vs.forEach(v -> logEntry.addRequestHeader(k, v)));

  attrs.putAttribute(LOG_ENTRY, logEntry.markStart());
}
 

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

/**
 * Validates that checksums should be enabled based on {@link ClientType} and the presence
 * or S3 specific headers.
 *
 * @param expectedClientType - The expected client type for enabling checksums
 * @param executionAttributes - {@link ExecutionAttributes} to determine the actual client type
 * @return If trailing checksums should be enabled for this request.
 */
public static boolean shouldRecordChecksum(SdkRequest sdkRequest,
                                           ClientType expectedClientType,
                                           ExecutionAttributes executionAttributes,
                                           SdkHttpRequest httpRequest) {
    if (!(sdkRequest instanceof PutObjectRequest)) {
        return false;
    }

    ClientType actualClientType = executionAttributes.getAttribute(SdkExecutionAttribute.CLIENT_TYPE);

    if (!expectedClientType.equals(actualClientType)) {
        return false;
    }


    if (hasServerSideEncryptionHeader(httpRequest)) {
        return false;
    }

    return checksumEnabledPerConfig(executionAttributes);
}
 

/**
 * After a SDK request has been executed
 */
@Override public void afterExecution(
    Context.AfterExecution context,
    ExecutionAttributes executionAttributes
) {
  Span span = executionAttributes.getAttribute(SPAN);
  if (span == null) {
    // An evil interceptor deleted our attribute.
    return;
  }
  handler.handleReceive(
      new HttpClientResponse(context.httpRequest(), context.httpResponse(), null), span);
}
 

protected <B extends Aws4SignerParams.Builder> B extractSignerParams(B paramsBuilder,
                                                                     ExecutionAttributes executionAttributes) {
    paramsBuilder.awsCredentials(executionAttributes.getAttribute(AwsSignerExecutionAttribute.AWS_CREDENTIALS))
                 .signingName(executionAttributes.getAttribute(AwsSignerExecutionAttribute.SERVICE_SIGNING_NAME))
                 .signingRegion(executionAttributes.getAttribute(AwsSignerExecutionAttribute.SIGNING_REGION))
                 .timeOffset(executionAttributes.getAttribute(AwsSignerExecutionAttribute.TIME_OFFSET));

    if (executionAttributes.getAttribute(AwsSignerExecutionAttribute.SIGNER_DOUBLE_URL_ENCODE) != null) {
        paramsBuilder.doubleUrlEncode(executionAttributes.getAttribute(AwsSignerExecutionAttribute.SIGNER_DOUBLE_URL_ENCODE));
    }

    return paramsBuilder;
}
 

/**
 * After a SDK request has failed
 */
@Override public void onExecutionFailure(
    Context.FailedExecution context,
    ExecutionAttributes executionAttributes
) {
  Span span = executionAttributes.getAttribute(SPAN);
  if (span == null) {
    // An evil interceptor deleted our attribute.
    return;
  }
  handler.handleReceive(new HttpClientResponse(
      context.httpRequest().orElse(null),
      context.httpResponse().orElse(null),
      maybeError(context)), span);
}
 

/**
 * Check the {@link S3Configuration#checksumValidationEnabled()} to see if the checksum is enabled.
 *
 * @param executionAttributes the execution attributes
 * @return true if the trailing checksum is enabled in the config, false otherwise.
 */
private static boolean checksumEnabledPerConfig(ExecutionAttributes executionAttributes) {
    S3Configuration serviceConfiguration =
        (S3Configuration) executionAttributes.getAttribute(AwsSignerExecutionAttribute.SERVICE_CONFIG);

    return serviceConfiguration == null || serviceConfiguration.checksumValidationEnabled();
}
 

/**
 * Before sending an http request. Will be called multiple times in the case of retries.
 */
@Override public void beforeTransmission(Context.BeforeTransmission context,
    ExecutionAttributes executionAttributes) {
  Span span = executionAttributes.getAttribute(SPAN);
  if (span == null) {
    // An evil interceptor deleted our attribute.
    return;
  }
  span.annotate("ws");
}
 

@Override
public void afterExecution(Context.AfterExecution context, ExecutionAttributes executionAttributes) {
    try {
        String uri = context.httpRequest().getUri().toString();
        EntityId nodeAccountId = getNodeAccountId(uri);
        Instant startTime = executionAttributes.getAttribute(START_TIME);
        ResponseSizeSubscriber responseSizeSubscriber = executionAttributes.getAttribute(SIZE);

        String[] tags = {
                "action", getAction(uri),
                "method", context.httpRequest().method().name(),
                "nodeAccount", nodeAccountId.getEntityNum().toString(),
                "realm", nodeAccountId.getRealmNum().toString(),
                "shard", nodeAccountId.getShardNum().toString(),
                "status", String.valueOf(context.httpResponse().statusCode()),
                "type", getType(uri)
        };

        if (startTime != null) {
            requestMetric.tags(tags)
                    .register(meterRegistry)
                    .record(Duration.between(startTime, Instant.now()));
        }

        if (responseSizeSubscriber != null) {
            responseSizeMetric.tags(tags)
                    .register(meterRegistry)
                    .record(Double.valueOf(responseSizeSubscriber.getSize()));
        }
    } catch (Exception e) {
        log.warn("Unable to collect S3 metrics", e);
    }
}
 

@Override
public void onExecutionFailure(Context.FailedExecution context, ExecutionAttributes executionAttributes) {
    Subsegment subsegment = executionAttributes.getAttribute(entityKey);
    if (subsegment == null) {
        return;
    }

    populateSubsegmentException(subsegment, context);
    populateRequestId(subsegment, context);
    if (context.httpResponse().isPresent()) {
        populateSubsegmentWithResponse(subsegment, context.httpResponse().get());
    }

    getRecorder().endSubsegment(subsegment);
}
 

/**
 * After sending an http request. Will be called multiple times in the case of retries.
 */
@Override public void afterTransmission(Context.AfterTransmission context,
    ExecutionAttributes executionAttributes) {
  Span span = executionAttributes.getAttribute(SPAN);
  if (span == null) {
    // An evil interceptor deleted our attribute.
    return;
  }
  span.annotate("wr");
}
 

/**
 * If there is a retry, then {@code beforeTransmission} will be called with the previous
 * attributes. This method will look for an existing entry and write out the log message.
 * The log entry may not have been filled in with a status if no response was received,
 * e.g., a connection exception. Since we do not have access to the failure, the status
 * will get set to {@code unexpected_error}.
 */
private void logRetryAttempt(ExecutionAttributes attrs) {
  IpcLogEntry logEntry = attrs.getAttribute(LOG_ENTRY);
  if (logEntry != null) {
    if (!isStatusSet(attrs)) {
      logEntry.markEnd().withStatus(IpcStatus.unexpected_error);
    }
    logEntry.log();
  }
}
 

@Override
public void onExecutionFailure(final FailedExecution context, final ExecutionAttributes executionAttributes) {
  final Span span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
  if (span == null)
    return;

  executionAttributes.putAttribute(SPAN_ATTRIBUTE, null);
  OpenTracingApiUtil.setErrorTag(span, context.exception());
  span.finish();
}
 

@Override
public void afterMarshalling(final AfterMarshalling context, final ExecutionAttributes executionAttributes) {
  final Span span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
  final SdkHttpRequest httpRequest = context.httpRequest();

  span.setTag(Tags.HTTP_METHOD, httpRequest.method().name());
  span.setTag(Tags.HTTP_URL, httpRequest.getUri().toString());
  span.setTag(Tags.PEER_HOSTNAME, httpRequest.host());
  if (httpRequest.port() > 0)
    span.setTag(Tags.PEER_PORT, httpRequest.port());
}
 

/**
 * Retrieve the region configured on the client.
 */
private Region clientRegion(ExecutionAttributes executionAttributes) {
    return executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION);
}
 

private Duration getClockSkew(ExecutionAttributes executionAttributes) {
    Integer timeOffset = executionAttributes.getAttribute(SdkExecutionAttribute.TIME_OFFSET);
    return timeOffset == null ? null : Duration.ofSeconds(timeOffset);
}
 

private Duration getClockSkew(ExecutionAttributes executionAttributes) {
    Integer timeOffset = executionAttributes.getAttribute(SdkExecutionAttribute.TIME_OFFSET);
    return timeOffset == null ? null : Duration.ofSeconds(timeOffset);
}
 

@Test
public void modifyRequest_DoesNotOverwriteExistingAttributeValue() {

    ExecutionAttributes executionAttributes = new ExecutionAttributes();

    interceptor.modifyRequest(context(PutObjectRequest.builder().build()), executionAttributes);

    boolean newValue = !executionAttributes.getAttribute(ENABLE_CHUNKED_ENCODING);

    executionAttributes.putAttribute(ENABLE_CHUNKED_ENCODING, newValue);

    interceptor.modifyRequest(context(PutObjectRequest.builder().build()), executionAttributes);

    assertThat(executionAttributes.getAttribute(ENABLE_CHUNKED_ENCODING)).isEqualTo(newValue);
}
 

/**
 * For network errors there will not be a response so the status will not have been set. This
 * method looks for a flag in the attributes to see if we need to close off the log entry for
 * the attempt.
 */
private boolean isStatusSet(ExecutionAttributes attrs) {
  Boolean s = attrs.getAttribute(STATUS_IS_SET);
  return s != null && s;
}