|
15 | 15 |
|
16 | 16 | package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v2_2; |
17 | 17 |
|
| 18 | +import static io.opentelemetry.instrumentation.api.internal.ConfigPropertiesUtil.getBoolean; |
| 19 | +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v2_2.AwsExperimentalAttributes.AWS_AUTH_ACCESS_KEY; |
| 20 | +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v2_2.AwsExperimentalAttributes.AWS_AUTH_REGION; |
| 21 | +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v2_2.AwsExperimentalAttributes.GEN_AI_SYSTEM; |
| 22 | +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v2_2.AwsSdkRequestType.BEDROCKRUNTIME; |
| 23 | + |
| 24 | +import io.opentelemetry.api.trace.Span; |
| 25 | +import software.amazon.awssdk.auth.credentials.AwsCredentials; |
| 26 | +import software.amazon.awssdk.auth.signer.AwsSignerExecutionAttribute; |
| 27 | +import software.amazon.awssdk.core.SdkRequest; |
18 | 28 | import software.amazon.awssdk.core.SdkResponse; |
19 | 29 | import software.amazon.awssdk.core.interceptor.*; |
| 30 | +import software.amazon.awssdk.regions.Region; |
20 | 31 |
|
| 32 | +/** |
| 33 | + * This interceptor manages the AWS SDK requests and responses for attribute mapping. Here's the |
| 34 | + * flow: |
| 35 | + * |
| 36 | + * <p>1. Request Phase (beforeTransmission): |
| 37 | + * |
| 38 | + * <ul> |
| 39 | + * <li>Receives the SDK request before it's sent to AWS |
| 40 | + * <li>Creates an AwsSdkRequest object containing request metadata |
| 41 | + * <li>Stores this AwsSdkRequest in ExecutionAttributes using ADOT_AWS_SDK_REQUEST_ATTRIBUTE |
| 42 | + * <li>Maps request attributes to the current span |
| 43 | + * </ul> |
| 44 | + * |
| 45 | + * <p>2. Response Phase (modifyResponse): |
| 46 | + * |
| 47 | + * <ul> |
| 48 | + * <li>Retrieves the stored AwsSdkRequest from ExecutionAttributes via |
| 49 | + * ADOT_AWS_SDK_REQUEST_ATTRIBUTE |
| 50 | + * <li>Uses this request context to properly map response fields to the span |
| 51 | + * <li>Cleans up by removing the stored request from ExecutionAttributes |
| 52 | + * </ul> |
| 53 | + * |
| 54 | + * <p>The ExecutionAttributes object persists throughout the entire request lifecycle, allowing |
| 55 | + * correlation between the request and response phases. All ExecutionInterceptor's have access to |
| 56 | + * the ExecutionAttributes. |
| 57 | + * |
| 58 | + * @see <a |
| 59 | + * href="https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/core/interceptor/ExecutionInterceptor.html">reference</a> |
| 60 | + */ |
21 | 61 | public class AdotAwsSdkTracingExecutionInterceptor implements ExecutionInterceptor { |
22 | 62 |
|
23 | | - // This is the latest point we can obtain the Sdk Request after it is modified by the upstream |
24 | | - // TracingInterceptor. It ensures upstream handles the request and applies its changes first. |
| 63 | + private static final String GEN_AI_SYSTEM_BEDROCK = "aws.bedrock"; |
| 64 | + private static final ExecutionAttribute<AwsSdkRequest> ADOT_AWS_SDK_REQUEST_ATTRIBUTE = |
| 65 | + new ExecutionAttribute<>( |
| 66 | + AdotAwsSdkTracingExecutionInterceptor.class.getName() + ".AwsSdkRequest"); |
| 67 | + |
| 68 | + private final FieldMapper fieldMapper = new FieldMapper(); |
| 69 | + private final boolean captureExperimentalSpanAttributes = |
| 70 | + getBoolean("otel.instrumentation.aws-sdk.experimental-span-attributes", true); |
| 71 | + |
| 72 | + /** |
| 73 | + * This method coordinates the request attribute mapping process. This is the latest point we can |
| 74 | + * obtain the Sdk Request after it is modified by the upstream TracingInterceptor. It ensures |
| 75 | + * upstream handles the request and applies its changes to the span first. We use this hook to |
| 76 | + * extract the ADOT AWS attributes from the Sdk Request and map them to the span via the |
| 77 | + * FieldMapper. |
| 78 | + * |
| 79 | + * <p>Upstream's last Sdk Request modification: @see <a |
| 80 | + * href="https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/release/v2.11.x/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/TracingExecutionInterceptor.java#L237">reference</a> |
| 81 | + */ |
25 | 82 | @Override |
26 | 83 | public void beforeTransmission( |
27 | | - Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {} |
| 84 | + Context.BeforeTransmission context, ExecutionAttributes executionAttributes) { |
28 | 85 |
|
29 | | - // This is the latest point we can obtain the Sdk Response before span completion in upstream's |
30 | | - // afterExecution. This ensures we capture attributes from the final, fully modified response |
31 | | - // after all upstream interceptors have processed it. |
| 86 | + if (captureExperimentalSpanAttributes) { |
| 87 | + SdkRequest request = context.request(); |
| 88 | + Span currentSpan = Span.current(); |
| 89 | + |
| 90 | + try { |
| 91 | + if (request == null || currentSpan == null || !currentSpan.getSpanContext().isValid()) { |
| 92 | + return; |
| 93 | + } |
| 94 | + |
| 95 | + AwsCredentials credentials = |
| 96 | + executionAttributes.getAttribute(AwsSignerExecutionAttribute.AWS_CREDENTIALS); |
| 97 | + Region signingRegion = |
| 98 | + executionAttributes.getAttribute(AwsSignerExecutionAttribute.SIGNING_REGION); |
| 99 | + |
| 100 | + if (credentials != null) { |
| 101 | + String accessKeyId = credentials.accessKeyId(); |
| 102 | + if (accessKeyId != null) { |
| 103 | + currentSpan.setAttribute(AWS_AUTH_ACCESS_KEY, accessKeyId); |
| 104 | + } |
| 105 | + } |
| 106 | + |
| 107 | + if (signingRegion != null) { |
| 108 | + String region = signingRegion.toString(); |
| 109 | + currentSpan.setAttribute(AWS_AUTH_REGION, region); |
| 110 | + } |
| 111 | + |
| 112 | + AwsSdkRequest awsSdkRequest = AwsSdkRequest.ofSdkRequest(request); |
| 113 | + if (awsSdkRequest != null) { |
| 114 | + executionAttributes.putAttribute(ADOT_AWS_SDK_REQUEST_ATTRIBUTE, awsSdkRequest); |
| 115 | + fieldMapper.mapToAttributes(request, awsSdkRequest, currentSpan); |
| 116 | + if (awsSdkRequest.type() == BEDROCKRUNTIME) { |
| 117 | + currentSpan.setAttribute(GEN_AI_SYSTEM, GEN_AI_SYSTEM_BEDROCK); |
| 118 | + } |
| 119 | + } |
| 120 | + } catch (Throwable throwable) { |
| 121 | + // ignore |
| 122 | + } |
| 123 | + } |
| 124 | + } |
| 125 | + |
| 126 | + /** |
| 127 | + * This method coordinates the response attribute mapping process. This is the latest point we can |
| 128 | + * obtain the Sdk Response before span completion in upstream's afterExecution. This ensures we |
| 129 | + * capture attributes from the final, fully modified response after all upstream interceptors have |
| 130 | + * processed it. We use this hook to extract the ADOT AWS attributes from the Sdk Response and map |
| 131 | + * them to the span via the FieldMapper. |
| 132 | + * |
| 133 | + * <p>Upstream's last Sdk Response modification before span closure: @see <a |
| 134 | + * href="https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/release/v2.11.x/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/TracingExecutionInterceptor.java#L348">reference</a> |
| 135 | + */ |
32 | 136 | @Override |
33 | 137 | public SdkResponse modifyResponse( |
34 | 138 | Context.ModifyResponse context, ExecutionAttributes executionAttributes) { |
35 | 139 |
|
| 140 | + if (captureExperimentalSpanAttributes) { |
| 141 | + Span currentSpan = Span.current(); |
| 142 | + AwsSdkRequest sdkRequest = executionAttributes.getAttribute(ADOT_AWS_SDK_REQUEST_ATTRIBUTE); |
| 143 | + |
| 144 | + if (sdkRequest != null) { |
| 145 | + fieldMapper.mapToAttributes(context.response(), sdkRequest, currentSpan); |
| 146 | + executionAttributes.putAttribute(ADOT_AWS_SDK_REQUEST_ATTRIBUTE, null); |
| 147 | + } |
| 148 | + } |
| 149 | + |
36 | 150 | return context.response(); |
37 | 151 | } |
38 | 152 | } |
0 commit comments