feat: AWS SDK v1 Step Functions auto-instrumentation support

instrumentation/aws-sdk/aws-sdk-1.11/javaagent/build.gradle.kts

@@ -41,12 +41,13 @@ dependencies {
 
   library("com.amazonaws:aws-java-sdk-core:1.11.0")
 
-  testLibrary("com.amazonaws:aws-java-sdk-s3:1.11.106")
-  testLibrary("com.amazonaws:aws-java-sdk-rds:1.11.106")
+  testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
   testLibrary("com.amazonaws:aws-java-sdk-ec2:1.11.106")
   testLibrary("com.amazonaws:aws-java-sdk-kinesis:1.11.106")
-  testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
+  testLibrary("com.amazonaws:aws-java-sdk-rds:1.11.106")
+  testLibrary("com.amazonaws:aws-java-sdk-s3:1.11.106")
   testLibrary("com.amazonaws:aws-java-sdk-sns:1.11.106")
+  testLibrary("com.amazonaws:aws-java-sdk-stepfunctions:1.11.106")
 
   testImplementation(project(":instrumentation:aws-sdk:aws-sdk-1.11:testing"))
 

instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/awssdk/v1_11/StepFunctionsClientTest.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.awssdk.v1_11;
+
+import com.amazonaws.services.stepfunctions.AWSStepFunctionsClientBuilder;
+import io.opentelemetry.instrumentation.awssdk.v1_11.AbstractStepFunctionsClientTest;
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+class StepFunctionsClientTest extends AbstractStepFunctionsClientTest {
+  @RegisterExtension
+  static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
+
+  @Override
+  protected InstrumentationExtension testing() {
+    return testing;
+  }
+
+  @Override
+  public AWSStepFunctionsClientBuilder configureClient(
+      AWSStepFunctionsClientBuilder clientBuilder) {
+    return clientBuilder;
+  }
+}

instrumentation/aws-sdk/aws-sdk-1.11/library/build.gradle.kts

@@ -11,12 +11,13 @@ dependencies {
 
   testImplementation(project(":instrumentation:aws-sdk:aws-sdk-1.11:testing"))
 
-  testLibrary("com.amazonaws:aws-java-sdk-s3:1.11.106")
-  testLibrary("com.amazonaws:aws-java-sdk-rds:1.11.106")
+  testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
   testLibrary("com.amazonaws:aws-java-sdk-ec2:1.11.106")
   testLibrary("com.amazonaws:aws-java-sdk-kinesis:1.11.106")
-  testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
+  testLibrary("com.amazonaws:aws-java-sdk-rds:1.11.106")
+  testLibrary("com.amazonaws:aws-java-sdk-s3:1.11.106")
   testLibrary("com.amazonaws:aws-java-sdk-sns:1.11.106")
+  testLibrary("com.amazonaws:aws-java-sdk-stepfunctions:1.11.106")
 
   // last version that does not use json protocol
   latestDepTestLibrary("com.amazonaws:aws-java-sdk-sqs:1.12.583") // documented limitation

instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkAttributesExtractor.java

@@ -16,12 +16,19 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
 import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
+import java.util.function.Function;
 import javax.annotation.Nullable;
 
 class AwsSdkAttributesExtractor implements AttributesExtractor<Request<?>, Response<?>> {
   private static final boolean CAN_GET_RESPONSE_METADATA = canGetResponseMetadata();
   private static final AttributeKey<String> AWS_REQUEST_ID = stringKey("aws.request_id");
 
+  // Copied from AwsIncubatingAttributes
+  private static final AttributeKey<String> AWS_STEP_FUNCTIONS_ACTIVITY_ARN =
+      stringKey("aws.step_functions.activity.arn");
+  private static final AttributeKey<String> AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN =
+      stringKey("aws.step_functions.state_machine.arn");
+
   // AmazonWebServiceResult is only available in v1.11.33 and later
   private static boolean canGetResponseMetadata() {
     try {
@@ -34,7 +41,19 @@ private static boolean canGetResponseMetadata() {
   }
 
   @Override
-  public void onStart(AttributesBuilder attributes, Context parentContext, Request<?> request) {}
+  public void onStart(AttributesBuilder attributes, Context parentContext, Request<?> request) {
+    Object originalRequest = request.getOriginalRequest();
+    setAttribute(
+        attributes,
+        AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN,
+        originalRequest,
+        RequestAccess::getStateMachineArn);
+    setAttribute(
+        attributes,
+        AWS_STEP_FUNCTIONS_ACTIVITY_ARN,
+        originalRequest,
+        RequestAccess::getStepFunctionsActivityArn);
+  }
 
   @Override
   public void onEnd(
@@ -43,8 +62,21 @@ public void onEnd(
       Request<?> request,
       @Nullable Response<?> response,
       @Nullable Throwable error) {
-    ResponseMetadata responseMetadata = getResponseMetadata(response);
+    if (response != null) {
+      Object awsResp = response.getAwsResponse();
+      setAttribute(
+          attributes,
+          AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN,
+          awsResp,
+          RequestAccess::getStateMachineArn);
+      setAttribute(
+          attributes,
+          AWS_STEP_FUNCTIONS_ACTIVITY_ARN,
+          awsResp,
+          RequestAccess::getStepFunctionsActivityArn);
+    }
 
+    ResponseMetadata responseMetadata = getResponseMetadata(response);
     if (responseMetadata != null) {
       String requestId = responseMetadata.getRequestId();
       if (requestId != null) {
@@ -63,4 +95,15 @@ private static ResponseMetadata getResponseMetadata(Response<?> response) {
     }
     return null;
   }
+
+  public static void setAttribute(
+      AttributesBuilder attributes,
+      AttributeKey<String> key,
+      Object carrier,
+      Function<Object, String> getter) {
+    String value = getter.apply(carrier);
+    if (value != null) {
+      attributes.put(key, value);
+    }
+  }
 }

instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/RequestAccess.java

@@ -11,6 +11,8 @@
 import javax.annotation.Nullable;
 
 final class RequestAccess {
+  private static final String STEP_FUNCTIONS_REQUEST_CLASS_PREFIX =
+      "com.amazonaws.services.stepfunctions.model.";
 
   private static final ClassValue<RequestAccess> REQUEST_ACCESSORS =
       new ClassValue<RequestAccess>() {
@@ -20,6 +22,24 @@ protected RequestAccess computeValue(Class<?> type) {
         }
       };
 
+  @Nullable
+  static String getStepFunctionsActivityArn(Object request) {
+    if (request == null) {
+      return null;
+    }
+    RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
+    return invokeOrNull(access.getStepFunctionsActivityArn, request);
+  }
+
+  @Nullable
+  static String getStateMachineArn(Object request) {
+    if (request == null) {
+      return null;
+    }
+    RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
+    return invokeOrNull(access.getStateMachineArn, request);
+  }
+
   @Nullable
   static String getBucketName(Object request) {
     RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
@@ -81,6 +101,8 @@ private static String invokeOrNull(@Nullable MethodHandle method, Object obj) {
   @Nullable private final MethodHandle getTableName;
   @Nullable private final MethodHandle getTopicArn;
   @Nullable private final MethodHandle getTargetArn;
+  @Nullable private final MethodHandle getStateMachineArn;
+  @Nullable private final MethodHandle getStepFunctionsActivityArn;
 
   private RequestAccess(Class<?> clz) {
     getBucketName = findAccessorOrNull(clz, "getBucketName");
@@ -90,6 +112,9 @@ private RequestAccess(Class<?> clz) {
     getTableName = findAccessorOrNull(clz, "getTableName");
     getTopicArn = findAccessorOrNull(clz, "getTopicArn");
     getTargetArn = findAccessorOrNull(clz, "getTargetArn");
+    boolean isStepFunction = clz.getName().startsWith(STEP_FUNCTIONS_REQUEST_CLASS_PREFIX);
+    getStateMachineArn = isStepFunction ? findAccessorOrNull(clz, "getStateMachineArn") : null;
+    getStepFunctionsActivityArn = isStepFunction ? findAccessorOrNull(clz, "getActivityArn") : null;
   }
 
   @Nullable

instrumentation/aws-sdk/aws-sdk-1.11/library/src/test/java/io/opentelemetry/instrumentation/awssdk/v1_11/StepFunctionsClientTest.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v1_11;
+
+import com.amazonaws.services.stepfunctions.AWSStepFunctionsClientBuilder;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.LibraryInstrumentationExtension;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+class StepFunctionsClientTest extends AbstractStepFunctionsClientTest {
+  @RegisterExtension
+  private static final InstrumentationExtension testing = LibraryInstrumentationExtension.create();
+
+  @Override
+  protected InstrumentationExtension testing() {
+    return testing;
+  }
+
+  @Override
+  public AWSStepFunctionsClientBuilder configureClient(
+      AWSStepFunctionsClientBuilder clientBuilder) {
+    return clientBuilder.withRequestHandlers(
+        AwsSdkTelemetry.builder(testing().getOpenTelemetry())
+            .setCaptureExperimentalSpanAttributes(true)
+            .build()
+            .newRequestHandler());
+  }
+}

instrumentation/aws-sdk/aws-sdk-1.11/testing/build.gradle.kts

@@ -7,13 +7,14 @@ dependencies {
 
   api("com.amazonaws:aws-java-sdk-core:1.11.0")
 
-  compileOnly("com.amazonaws:aws-java-sdk-s3:1.11.106")
-  compileOnly("com.amazonaws:aws-java-sdk-rds:1.11.106")
+  compileOnly("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
   compileOnly("com.amazonaws:aws-java-sdk-ec2:1.11.106")
   compileOnly("com.amazonaws:aws-java-sdk-kinesis:1.11.106")
-  compileOnly("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
+  compileOnly("com.amazonaws:aws-java-sdk-rds:1.11.106")
+  compileOnly("com.amazonaws:aws-java-sdk-s3:1.11.106")
   compileOnly("com.amazonaws:aws-java-sdk-sns:1.11.106")
   compileOnly("com.amazonaws:aws-java-sdk-sqs:1.11.106")
+  compileOnly("com.amazonaws:aws-java-sdk-stepfunctions:1.11.106")
 
   // needed for SQS - using emq directly as localstack references emq v0.15.7 ie WITHOUT AWS trace header propagation
   implementation("org.elasticmq:elasticmq-rest-sqs_2.13")

instrumentation/aws-sdk/aws-sdk-1.11/testing/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AbstractStepFunctionsClientTest.java

@@ -0,0 +1,78 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v1_11;
+
+import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
+import static io.opentelemetry.semconv.incubating.AwsIncubatingAttributes.AWS_STEP_FUNCTIONS_ACTIVITY_ARN;
+import static io.opentelemetry.semconv.incubating.AwsIncubatingAttributes.AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN;
+import static java.util.Collections.singletonList;
+
+import com.amazonaws.services.stepfunctions.AWSStepFunctions;
+import com.amazonaws.services.stepfunctions.AWSStepFunctionsClientBuilder;
+import com.amazonaws.services.stepfunctions.model.DescribeActivityRequest;
+import com.amazonaws.services.stepfunctions.model.DescribeStateMachineRequest;
+import io.opentelemetry.sdk.testing.assertj.AttributeAssertion;
+import io.opentelemetry.testing.internal.armeria.common.HttpResponse;
+import io.opentelemetry.testing.internal.armeria.common.HttpStatus;
+import io.opentelemetry.testing.internal.armeria.common.MediaType;
+import java.util.List;
+import java.util.function.Function;
+import java.util.stream.Stream;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+public abstract class AbstractStepFunctionsClientTest extends AbstractBaseAwsClientTest {
+
+  public abstract AWSStepFunctionsClientBuilder configureClient(
+      AWSStepFunctionsClientBuilder client);
+
+  @Override
+  protected boolean hasRequestId() {
+    return false;
+  }
+
+  @ParameterizedTest
+  @MethodSource("provideArguments")
+  public void testSendRequestWithMockedResponse(
+      String operation,
+      List<AttributeAssertion> additionalAttributes,
+      Function<AWSStepFunctions, Object> call)
+      throws Exception {
+
+    AWSStepFunctionsClientBuilder clientBuilder = AWSStepFunctionsClientBuilder.standard();
+
+    AWSStepFunctions client =
+        configureClient(clientBuilder)
+            .withEndpointConfiguration(endpoint)
+            .withCredentials(credentialsProvider)
+            .build();
+
+    server.enqueue(HttpResponse.of(HttpStatus.OK, MediaType.PLAIN_TEXT_UTF_8, ""));
+
+    Object response = call.apply(client);
+    assertRequestWithMockedResponse(
+        response, client, "AWSStepFunctions", operation, "POST", additionalAttributes);
+  }
+
+  private static Stream<Arguments> provideArguments() {
+    return Stream.of(
+        Arguments.of(
+            "DescribeStateMachine",
+            singletonList(equalTo(AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN, "stateMachineArn")),
+            (Function<AWSStepFunctions, Object>)
+                c ->
+                    c.describeStateMachine(
+                        new DescribeStateMachineRequest().withStateMachineArn("stateMachineArn"))),
+        Arguments.of(
+            "DescribeActivity",
+            singletonList(equalTo(AWS_STEP_FUNCTIONS_ACTIVITY_ARN, "activityArn")),
+            (Function<AWSStepFunctions, Object>)
+                c ->
+                    c.describeActivity(
+                        new DescribeActivityRequest().withActivityArn("activityArn"))));
+  }
+}