AWS SDK v1.11 Patch Migration

instrumentation/aws-sdk/README.md

@@ -27,6 +27,21 @@ The aws-sdk instrumentation is an SPI-based implementation that extends the upst
    - Finds ADOT’s **AdotAwsSdkInstrumentationModule**
    - Registers **AdotAwsSdkTracingExecutionInterceptor** (order > 0)
 
+#### _Note on Attribute Collection:_
+AWS SDK v1.11 and v2.2 handle attribute collection differently:
+
+**V1.11:**
+- Maintains a separate AttributesBuilder during request/response lifecycle
+- Collects ADOT-specific attributes alongside upstream processing without interference
+- Injects collected attributes into span at the end of the request and response lifecycle hooks
+
+
+**V2.2:**
+- FieldMapper directly modifies spans during request/response processing
+- Attributes are added to spans immediately when discovered
+- Direct integration with span lifecycle
+
+This architectural difference exists due to upstream AWS SDK injecting attributes into spans differently for v1.11 and v2.2
 ### AWS SDK v1 Instrumentation Summary
 The AdotAwsSdkInstrumentationModule uses the instrumentation (specified in AdotAwsClientInstrumentation) to register the AdotAwsSdkTracingRequestHandler through `typeInstrumentations`.
 
@@ -61,6 +76,28 @@ _**Important Notes:**_
 - The upstream interceptor closes the span in [afterResponse](https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java#L116) and/or [afterError](https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java#L131). These hooks are inaccessible for span modification.
   `afterAttempt` is our final hook point, giving us access to both the fully processed response and active span.
 
+**High-Level Sequence Diagram:**
+![img.png](sequence-diagram-1.11.png)
+
+_Class Functionalities:_
+- `AdotAwsSdkTracingRequestHandler`
+  - Hooks into AWS SDK request/response lifecycle 
+  - Adds ADOT-specific attributes to spans extracted by AwsSdkExperimentalAttributesExtractor
+- `AwsSdkExperimentalAttributesExtractor`
+  - Extracts attributes from AWS requests/responses and enriches spans
+  - Uses RequestAccess to get field values
+  - Special handling for Bedrock services
+- `RequestAccess`
+  - Provides access to AWS SDK object fields 
+  - Caches method handles for performance 
+  - Uses BedrockJsonParser for parsing LLM payloads
+- `BedrockJsonParser`
+  - Custom JSON parser for Bedrock payloads
+  - Handles different LLM model formats
+- `AwsBedrockResourceType`
+  - Maps Bedrock class names to resource types
+  - Provides attribute keys and accessors for each type
+
 ### AWS SDK v2 Instrumentation Summary
 
 **AdotAwsSdkInstrumentationModule**
@@ -83,4 +120,14 @@ _**Important Note:**_
 The upstream interceptor closes the span in `afterExecution`. That hook is inaccessible for span modification.
 `modifyResponse` is our final hook point, giving us access to both the fully processed response and active span.
 
+### Commands for Running Groovy Tests
+
+To run the BedrockJsonParserTest for aws-sdk v1.11:
+````
+./gradlew :instrumentation:aws-sdk:test --tests "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.BedrockJsonParserTest"
+````
 
+To run the BedrockJsonParserTest for aws-sdk v2.2:
+````
+./gradlew :instrumentation:aws-sdk:test --tests "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v2_2.BedrockJsonParserTest"
+````

instrumentation/aws-sdk/build.gradle.kts

@@ -16,17 +16,22 @@
 plugins {
   java
   id("com.gradleup.shadow")
+  id("groovy")
 }
 
 base.archivesBaseName = "aws-instrumentation-aws-sdk"
 
 dependencies {
+  compileOnly("com.google.code.findbugs:jsr305:3.0.2")
   compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-extension-api")
+  compileOnly("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api")
   compileOnly("com.amazonaws:aws-java-sdk-core:1.11.0")
   compileOnly("software.amazon.awssdk:aws-core:2.2.0")
   compileOnly("net.bytebuddy:byte-buddy")
 
   testImplementation("com.amazonaws:aws-java-sdk-core:1.11.0")
   testImplementation("io.opentelemetry.javaagent:opentelemetry-javaagent-extension-api")
   testImplementation("org.mockito:mockito-core:5.14.2")
+  testImplementation("com.google.guava:guava")
+  testImplementation("io.opentelemetry.javaagent:opentelemetry-testing-common")
 }

instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AdotAwsSdkInstrumentationModule.java

@@ -16,6 +16,7 @@
 package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11;
 
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
+import static net.bytebuddy.matcher.ElementMatchers.*;
 
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
@@ -46,7 +47,17 @@ public int order() {
   @Override
   public List<String> getAdditionalHelperClassNames() {
     return Arrays.asList(
-        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AdotAwsSdkTracingRequestHandler");
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AdotAwsSdkTracingRequestHandler",
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsSdkExperimentalAttributesExtractor",
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes",
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsBedrockResourceType",
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsBedrockResourceType$AwsBedrockResourceTypeMap",
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.BedrockJsonParser",
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.BedrockJsonParser$JsonParser",
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.BedrockJsonParser$LlmJson",
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.BedrockJsonParser$JsonPathResolver",
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.RequestAccess",
+        "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.RequestAccess$1");
   }
 
   @Override

instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AdotAwsSdkTracingRequestHandler.java

@@ -16,8 +16,13 @@
 package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11;
 
 import com.amazonaws.Request;
+import com.amazonaws.Response;
 import com.amazonaws.handlers.HandlerAfterAttemptContext;
 import com.amazonaws.handlers.RequestHandler2;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
 
 /**
  * Based on OpenTelemetry Java Instrumentation's AWS SDK v1.11 TracingRequestHandler
@@ -27,8 +32,11 @@
  * href="https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/release/v2.11.x/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java">...</a>
  */
 public class AdotAwsSdkTracingRequestHandler extends RequestHandler2 {
+  private final AwsSdkExperimentalAttributesExtractor experimentalAttributesExtractor;
 
-  public AdotAwsSdkTracingRequestHandler() {}
+  public AdotAwsSdkTracingRequestHandler() {
+    this.experimentalAttributesExtractor = new AwsSdkExperimentalAttributesExtractor();
+  }
 
   /**
    * This is the latest point we can obtain the Sdk Request after it is modified by the upstream
@@ -38,7 +46,18 @@ public AdotAwsSdkTracingRequestHandler() {}
    * href="https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java#L58">reference</a>
    */
   @Override
-  public void beforeRequest(Request<?> request) {}
+  public void beforeRequest(Request<?> request) {
+    Span currentSpan = Span.current();
+
+    if (currentSpan != null && currentSpan.getSpanContext().isValid()) {
+      AttributesBuilder attributes = Attributes.builder();
+      experimentalAttributesExtractor.onStart(attributes, Context.current(), request);
+
+      attributes
+          .build()
+          .forEach((key, value) -> currentSpan.setAttribute(key.getKey(), value.toString()));
+    }
+  }
 
   /**
    * This is the latest point to access the sdk response before the span closes in the upstream
@@ -52,5 +71,21 @@ public void beforeRequest(Request<?> request) {}
    *     href="https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java#L131">reference</a>
    */
   @Override
-  public void afterAttempt(HandlerAfterAttemptContext context) {}
+  public void afterAttempt(HandlerAfterAttemptContext context) {
+    Span currentSpan = Span.current();
+
+    if (currentSpan != null && currentSpan.getSpanContext().isValid()) {
+      Request<?> request = context.getRequest();
+      Response<?> response = context.getResponse();
+      Exception exception = context.getException();
+
+      AttributesBuilder attributes = Attributes.builder();
+      experimentalAttributesExtractor.onEnd(
+          attributes, Context.current(), request, response, exception);
+
+      attributes
+          .build()
+          .forEach((key, value) -> currentSpan.setAttribute(key.getKey(), value.toString()));
+    }
+  }
 }

instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AwsBedrockResourceType.java

@@ -0,0 +1,143 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11;
+
+import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_AGENT_ID;
+import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_DATA_SOURCE_ID;
+import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_KNOWLEDGE_BASE_ID;
+
+import io.opentelemetry.api.common.AttributeKey;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+
+enum AwsBedrockResourceType {
+  AGENT_TYPE(AWS_AGENT_ID, RequestAccess::getAgentId),
+  DATA_SOURCE_TYPE(AWS_DATA_SOURCE_ID, RequestAccess::getDataSourceId),
+  KNOWLEDGE_BASE_TYPE(AWS_KNOWLEDGE_BASE_ID, RequestAccess::getKnowledgeBaseId);
+
+  @SuppressWarnings("ImmutableEnumChecker")
+  private final AttributeKey<String> keyAttribute;
+
+  @SuppressWarnings("ImmutableEnumChecker")
+  private final Function<Object, String> attributeValueAccessor;
+
+  AwsBedrockResourceType(
+      AttributeKey<String> keyAttribute, Function<Object, String> attributeValueAccessor) {
+    this.keyAttribute = keyAttribute;
+    this.attributeValueAccessor = attributeValueAccessor;
+  }
+
+  public AttributeKey<String> getKeyAttribute() {
+    return keyAttribute;
+  }
+
+  public Function<Object, String> getAttributeValueAccessor() {
+    return attributeValueAccessor;
+  }
+
+  public static AwsBedrockResourceType getRequestType(String requestClass) {
+    return AwsBedrockResourceTypeMap.BEDROCK_REQUEST_MAP.get(requestClass);
+  }
+
+  public static AwsBedrockResourceType getResponseType(String responseClass) {
+    return AwsBedrockResourceTypeMap.BEDROCK_RESPONSE_MAP.get(responseClass);
+  }
+
+  private static class AwsBedrockResourceTypeMap {
+    private static final Map<String, AwsBedrockResourceType> BEDROCK_REQUEST_MAP = new HashMap<>();
+    private static final Map<String, AwsBedrockResourceType> BEDROCK_RESPONSE_MAP = new HashMap<>();
+
+    // Bedrock request/response mapping
+    // We only support operations that are related to the resource and where the context contains
+    // the AgentID/DataSourceID/KnowledgeBaseID.
+    // AgentID
+    private static final List<String> agentRequestClasses =
+        Arrays.asList(
+            "CreateAgentActionGroupRequest",
+            "CreateAgentAliasRequest",
+            "DeleteAgentActionGroupRequest",
+            "DeleteAgentAliasRequest",
+            "DeleteAgentRequest",
+            "DeleteAgentVersionRequest",
+            "GetAgentActionGroupRequest",
+            "GetAgentAliasRequest",
+            "GetAgentRequest",
+            "GetAgentVersionRequest",
+            "ListAgentActionGroupsRequest",
+            "ListAgentAliasesRequest",
+            "ListAgentKnowledgeBasesRequest",
+            "ListAgentVersionsRequest",
+            "PrepareAgentRequest",
+            "UpdateAgentActionGroupRequest",
+            "UpdateAgentAliasRequest",
+            "UpdateAgentRequest");
+    private static final List<String> agentResponseClasses =
+        Arrays.asList(
+            "DeleteAgentAliasResult",
+            "DeleteAgentResult",
+            "DeleteAgentVersionResult",
+            "PrepareAgentResult");
+    // DataSourceID
+    private static final List<String> dataSourceRequestClasses =
+        Arrays.asList("DeleteDataSourceRequest", "GetDataSourceRequest", "UpdateDataSourceRequest");
+    private static final List<String> dataSourceResponseClasses =
+        Arrays.asList("DeleteDataSourceResult");
+    // KnowledgeBaseID
+    private static final List<String> knowledgeBaseRequestClasses =
+        Arrays.asList(
+            "AssociateAgentKnowledgeBaseRequest",
+            "CreateDataSourceRequest",
+            "DeleteKnowledgeBaseRequest",
+            "DisassociateAgentKnowledgeBaseRequest",
+            "GetAgentKnowledgeBaseRequest",
+            "GetKnowledgeBaseRequest",
+            "ListDataSourcesRequest",
+            "UpdateAgentKnowledgeBaseRequest");
+    private static final List<String> knowledgeBaseResponseClasses =
+        Arrays.asList("DeleteKnowledgeBaseResult");
+
+    private AwsBedrockResourceTypeMap() {}
+
+    static {
+      // Populate the BEDROCK_REQUEST_MAP
+      for (String agentRequestClass : agentRequestClasses) {
+        BEDROCK_REQUEST_MAP.put(agentRequestClass, AwsBedrockResourceType.AGENT_TYPE);
+      }
+      for (String dataSourceRequestClass : dataSourceRequestClasses) {
+        BEDROCK_REQUEST_MAP.put(dataSourceRequestClass, AwsBedrockResourceType.DATA_SOURCE_TYPE);
+      }
+      for (String knowledgeBaseRequestClass : knowledgeBaseRequestClasses) {
+        BEDROCK_REQUEST_MAP.put(
+            knowledgeBaseRequestClass, AwsBedrockResourceType.KNOWLEDGE_BASE_TYPE);
+      }
+
+      // Populate the BEDROCK_RESPONSE_MAP
+      for (String agentResponseClass : agentResponseClasses) {
+        BEDROCK_REQUEST_MAP.put(agentResponseClass, AwsBedrockResourceType.AGENT_TYPE);
+      }
+      for (String dataSourceResponseClass : dataSourceResponseClasses) {
+        BEDROCK_REQUEST_MAP.put(dataSourceResponseClass, AwsBedrockResourceType.DATA_SOURCE_TYPE);
+      }
+      for (String knowledgeBaseResponseClass : knowledgeBaseResponseClasses) {
+        BEDROCK_REQUEST_MAP.put(
+            knowledgeBaseResponseClass, AwsBedrockResourceType.KNOWLEDGE_BASE_TYPE);
+      }
+    }
+  }
+}