Skip to content

Commit c726d32

Browse files
lukeina2zlaurit
andauthored
feat: AWS SDK v1 Step Functions auto-instrumentation support (#14003)
Co-authored-by: Lauri Tulmin <[email protected]>
1 parent 4a2b49f commit c726d32

File tree

8 files changed

+219
-11
lines changed

8 files changed

+219
-11
lines changed

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,13 @@ dependencies {
4141

4242
library("com.amazonaws:aws-java-sdk-core:1.11.0")
4343

44-
testLibrary("com.amazonaws:aws-java-sdk-s3:1.11.106")
45-
testLibrary("com.amazonaws:aws-java-sdk-rds:1.11.106")
44+
testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
4645
testLibrary("com.amazonaws:aws-java-sdk-ec2:1.11.106")
4746
testLibrary("com.amazonaws:aws-java-sdk-kinesis:1.11.106")
48-
testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
47+
testLibrary("com.amazonaws:aws-java-sdk-rds:1.11.106")
48+
testLibrary("com.amazonaws:aws-java-sdk-s3:1.11.106")
4949
testLibrary("com.amazonaws:aws-java-sdk-sns:1.11.106")
50+
testLibrary("com.amazonaws:aws-java-sdk-stepfunctions:1.11.106")
5051

5152
testImplementation(project(":instrumentation:aws-sdk:aws-sdk-1.11:testing"))
5253

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.awssdk.v1_11;
7+
8+
import com.amazonaws.services.stepfunctions.AWSStepFunctionsClientBuilder;
9+
import io.opentelemetry.instrumentation.awssdk.v1_11.AbstractStepFunctionsClientTest;
10+
import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
11+
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
12+
import org.junit.jupiter.api.extension.RegisterExtension;
13+
14+
class StepFunctionsClientTest extends AbstractStepFunctionsClientTest {
15+
@RegisterExtension
16+
static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
17+
18+
@Override
19+
protected InstrumentationExtension testing() {
20+
return testing;
21+
}
22+
23+
@Override
24+
public AWSStepFunctionsClientBuilder configureClient(
25+
AWSStepFunctionsClientBuilder clientBuilder) {
26+
return clientBuilder;
27+
}
28+
}

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,13 @@ dependencies {
1111

1212
testImplementation(project(":instrumentation:aws-sdk:aws-sdk-1.11:testing"))
1313

14-
testLibrary("com.amazonaws:aws-java-sdk-s3:1.11.106")
15-
testLibrary("com.amazonaws:aws-java-sdk-rds:1.11.106")
14+
testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
1615
testLibrary("com.amazonaws:aws-java-sdk-ec2:1.11.106")
1716
testLibrary("com.amazonaws:aws-java-sdk-kinesis:1.11.106")
18-
testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
17+
testLibrary("com.amazonaws:aws-java-sdk-rds:1.11.106")
18+
testLibrary("com.amazonaws:aws-java-sdk-s3:1.11.106")
1919
testLibrary("com.amazonaws:aws-java-sdk-sns:1.11.106")
20+
testLibrary("com.amazonaws:aws-java-sdk-stepfunctions:1.11.106")
2021

2122
// last version that does not use json protocol
2223
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

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,19 @@
1616
import io.opentelemetry.context.Context;
1717
import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
1818
import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
19+
import java.util.function.Function;
1920
import javax.annotation.Nullable;
2021

2122
class AwsSdkAttributesExtractor implements AttributesExtractor<Request<?>, Response<?>> {
2223
private static final boolean CAN_GET_RESPONSE_METADATA = canGetResponseMetadata();
2324
private static final AttributeKey<String> AWS_REQUEST_ID = stringKey("aws.request_id");
2425

26+
// Copied from AwsIncubatingAttributes
27+
private static final AttributeKey<String> AWS_STEP_FUNCTIONS_ACTIVITY_ARN =
28+
stringKey("aws.step_functions.activity.arn");
29+
private static final AttributeKey<String> AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN =
30+
stringKey("aws.step_functions.state_machine.arn");
31+
2532
// AmazonWebServiceResult is only available in v1.11.33 and later
2633
private static boolean canGetResponseMetadata() {
2734
try {
@@ -34,7 +41,19 @@ private static boolean canGetResponseMetadata() {
3441
}
3542

3643
@Override
37-
public void onStart(AttributesBuilder attributes, Context parentContext, Request<?> request) {}
44+
public void onStart(AttributesBuilder attributes, Context parentContext, Request<?> request) {
45+
Object originalRequest = request.getOriginalRequest();
46+
setAttribute(
47+
attributes,
48+
AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN,
49+
originalRequest,
50+
RequestAccess::getStateMachineArn);
51+
setAttribute(
52+
attributes,
53+
AWS_STEP_FUNCTIONS_ACTIVITY_ARN,
54+
originalRequest,
55+
RequestAccess::getStepFunctionsActivityArn);
56+
}
3857

3958
@Override
4059
public void onEnd(
@@ -43,8 +62,21 @@ public void onEnd(
4362
Request<?> request,
4463
@Nullable Response<?> response,
4564
@Nullable Throwable error) {
46-
ResponseMetadata responseMetadata = getResponseMetadata(response);
65+
if (response != null) {
66+
Object awsResp = response.getAwsResponse();
67+
setAttribute(
68+
attributes,
69+
AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN,
70+
awsResp,
71+
RequestAccess::getStateMachineArn);
72+
setAttribute(
73+
attributes,
74+
AWS_STEP_FUNCTIONS_ACTIVITY_ARN,
75+
awsResp,
76+
RequestAccess::getStepFunctionsActivityArn);
77+
}
4778

79+
ResponseMetadata responseMetadata = getResponseMetadata(response);
4880
if (responseMetadata != null) {
4981
String requestId = responseMetadata.getRequestId();
5082
if (requestId != null) {
@@ -63,4 +95,15 @@ private static ResponseMetadata getResponseMetadata(Response<?> response) {
6395
}
6496
return null;
6597
}
98+
99+
public static void setAttribute(
100+
AttributesBuilder attributes,
101+
AttributeKey<String> key,
102+
Object carrier,
103+
Function<Object, String> getter) {
104+
String value = getter.apply(carrier);
105+
if (value != null) {
106+
attributes.put(key, value);
107+
}
108+
}
66109
}

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
import javax.annotation.Nullable;
1212

1313
final class RequestAccess {
14+
private static final String STEP_FUNCTIONS_REQUEST_CLASS_PREFIX =
15+
"com.amazonaws.services.stepfunctions.model.";
1416

1517
private static final ClassValue<RequestAccess> REQUEST_ACCESSORS =
1618
new ClassValue<RequestAccess>() {
@@ -20,6 +22,24 @@ protected RequestAccess computeValue(Class<?> type) {
2022
}
2123
};
2224

25+
@Nullable
26+
static String getStepFunctionsActivityArn(Object request) {
27+
if (request == null) {
28+
return null;
29+
}
30+
RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
31+
return invokeOrNull(access.getStepFunctionsActivityArn, request);
32+
}
33+
34+
@Nullable
35+
static String getStateMachineArn(Object request) {
36+
if (request == null) {
37+
return null;
38+
}
39+
RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
40+
return invokeOrNull(access.getStateMachineArn, request);
41+
}
42+
2343
@Nullable
2444
static String getBucketName(Object request) {
2545
RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
@@ -81,6 +101,8 @@ private static String invokeOrNull(@Nullable MethodHandle method, Object obj) {
81101
@Nullable private final MethodHandle getTableName;
82102
@Nullable private final MethodHandle getTopicArn;
83103
@Nullable private final MethodHandle getTargetArn;
104+
@Nullable private final MethodHandle getStateMachineArn;
105+
@Nullable private final MethodHandle getStepFunctionsActivityArn;
84106

85107
private RequestAccess(Class<?> clz) {
86108
getBucketName = findAccessorOrNull(clz, "getBucketName");
@@ -90,6 +112,9 @@ private RequestAccess(Class<?> clz) {
90112
getTableName = findAccessorOrNull(clz, "getTableName");
91113
getTopicArn = findAccessorOrNull(clz, "getTopicArn");
92114
getTargetArn = findAccessorOrNull(clz, "getTargetArn");
115+
boolean isStepFunction = clz.getName().startsWith(STEP_FUNCTIONS_REQUEST_CLASS_PREFIX);
116+
getStateMachineArn = isStepFunction ? findAccessorOrNull(clz, "getStateMachineArn") : null;
117+
getStepFunctionsActivityArn = isStepFunction ? findAccessorOrNull(clz, "getActivityArn") : null;
93118
}
94119

95120
@Nullable
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.awssdk.v1_11;
7+
8+
import com.amazonaws.services.stepfunctions.AWSStepFunctionsClientBuilder;
9+
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
10+
import io.opentelemetry.instrumentation.testing.junit.LibraryInstrumentationExtension;
11+
import org.junit.jupiter.api.extension.RegisterExtension;
12+
13+
class StepFunctionsClientTest extends AbstractStepFunctionsClientTest {
14+
@RegisterExtension
15+
private static final InstrumentationExtension testing = LibraryInstrumentationExtension.create();
16+
17+
@Override
18+
protected InstrumentationExtension testing() {
19+
return testing;
20+
}
21+
22+
@Override
23+
public AWSStepFunctionsClientBuilder configureClient(
24+
AWSStepFunctionsClientBuilder clientBuilder) {
25+
return clientBuilder.withRequestHandlers(
26+
AwsSdkTelemetry.builder(testing().getOpenTelemetry())
27+
.setCaptureExperimentalSpanAttributes(true)
28+
.build()
29+
.newRequestHandler());
30+
}
31+
}

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,14 @@ dependencies {
77

88
api("com.amazonaws:aws-java-sdk-core:1.11.0")
99

10-
compileOnly("com.amazonaws:aws-java-sdk-s3:1.11.106")
11-
compileOnly("com.amazonaws:aws-java-sdk-rds:1.11.106")
10+
compileOnly("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
1211
compileOnly("com.amazonaws:aws-java-sdk-ec2:1.11.106")
1312
compileOnly("com.amazonaws:aws-java-sdk-kinesis:1.11.106")
14-
compileOnly("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
13+
compileOnly("com.amazonaws:aws-java-sdk-rds:1.11.106")
14+
compileOnly("com.amazonaws:aws-java-sdk-s3:1.11.106")
1515
compileOnly("com.amazonaws:aws-java-sdk-sns:1.11.106")
1616
compileOnly("com.amazonaws:aws-java-sdk-sqs:1.11.106")
17+
compileOnly("com.amazonaws:aws-java-sdk-stepfunctions:1.11.106")
1718

1819
// needed for SQS - using emq directly as localstack references emq v0.15.7 ie WITHOUT AWS trace header propagation
1920
implementation("org.elasticmq:elasticmq-rest-sqs_2.13")
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.awssdk.v1_11;
7+
8+
import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
9+
import static io.opentelemetry.semconv.incubating.AwsIncubatingAttributes.AWS_STEP_FUNCTIONS_ACTIVITY_ARN;
10+
import static io.opentelemetry.semconv.incubating.AwsIncubatingAttributes.AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN;
11+
import static java.util.Collections.singletonList;
12+
13+
import com.amazonaws.services.stepfunctions.AWSStepFunctions;
14+
import com.amazonaws.services.stepfunctions.AWSStepFunctionsClientBuilder;
15+
import com.amazonaws.services.stepfunctions.model.DescribeActivityRequest;
16+
import com.amazonaws.services.stepfunctions.model.DescribeStateMachineRequest;
17+
import io.opentelemetry.sdk.testing.assertj.AttributeAssertion;
18+
import io.opentelemetry.testing.internal.armeria.common.HttpResponse;
19+
import io.opentelemetry.testing.internal.armeria.common.HttpStatus;
20+
import io.opentelemetry.testing.internal.armeria.common.MediaType;
21+
import java.util.List;
22+
import java.util.function.Function;
23+
import java.util.stream.Stream;
24+
import org.junit.jupiter.params.ParameterizedTest;
25+
import org.junit.jupiter.params.provider.Arguments;
26+
import org.junit.jupiter.params.provider.MethodSource;
27+
28+
public abstract class AbstractStepFunctionsClientTest extends AbstractBaseAwsClientTest {
29+
30+
public abstract AWSStepFunctionsClientBuilder configureClient(
31+
AWSStepFunctionsClientBuilder client);
32+
33+
@Override
34+
protected boolean hasRequestId() {
35+
return false;
36+
}
37+
38+
@ParameterizedTest
39+
@MethodSource("provideArguments")
40+
public void testSendRequestWithMockedResponse(
41+
String operation,
42+
List<AttributeAssertion> additionalAttributes,
43+
Function<AWSStepFunctions, Object> call)
44+
throws Exception {
45+
46+
AWSStepFunctionsClientBuilder clientBuilder = AWSStepFunctionsClientBuilder.standard();
47+
48+
AWSStepFunctions client =
49+
configureClient(clientBuilder)
50+
.withEndpointConfiguration(endpoint)
51+
.withCredentials(credentialsProvider)
52+
.build();
53+
54+
server.enqueue(HttpResponse.of(HttpStatus.OK, MediaType.PLAIN_TEXT_UTF_8, ""));
55+
56+
Object response = call.apply(client);
57+
assertRequestWithMockedResponse(
58+
response, client, "AWSStepFunctions", operation, "POST", additionalAttributes);
59+
}
60+
61+
private static Stream<Arguments> provideArguments() {
62+
return Stream.of(
63+
Arguments.of(
64+
"DescribeStateMachine",
65+
singletonList(equalTo(AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN, "stateMachineArn")),
66+
(Function<AWSStepFunctions, Object>)
67+
c ->
68+
c.describeStateMachine(
69+
new DescribeStateMachineRequest().withStateMachineArn("stateMachineArn"))),
70+
Arguments.of(
71+
"DescribeActivity",
72+
singletonList(equalTo(AWS_STEP_FUNCTIONS_ACTIVITY_ARN, "activityArn")),
73+
(Function<AWSStepFunctions, Object>)
74+
c ->
75+
c.describeActivity(
76+
new DescribeActivityRequest().withActivityArn("activityArn"))));
77+
}
78+
}

0 commit comments

Comments
 (0)