Skip to content

feat: Add auto-instrumentation support in AWS V1 SDK for Step Functio… #11

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 4, 2024
Merged

feat: Add auto-instrumentation support in AWS V1 SDK for Step Functio… #11

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions ...rc/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsExperimentalAttributes.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,5 +29,20 @@ final class AwsExperimentalAttributes {
stringKey("gen_ai.request.model");
static final AttributeKey<String> AWS_BEDROCK_SYSTEM = stringKey("gen_ai.system");

static final AttributeKey<String> AWS_STATE_MACHINE_ARN =
stringKey("aws.stepfunctions.state_machine.arn");

static final AttributeKey<String> AWS_STEP_FUNCTIONS_ACTIVITY_ARN =
stringKey("aws.stepfunctions.activity.arn");

static final AttributeKey<String> AWS_SNS_TOPIC_ARN = stringKey("aws.sns.topic.arn");

static final AttributeKey<String> AWS_SECRET_ARN = stringKey("aws.secretsmanager.secret.arn");

static final AttributeKey<String> AWS_LAMBDA_NAME = stringKey("aws.lambda.function.name");
Copy link
Owner

@mxiamxia mxiamxia Oct 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we also get lambda function.arn from the response?

Copy link
Author

@yiyuan-he yiyuan-he Oct 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure I can try to implement this

Copy link
Author

@yiyuan-he yiyuan-he Oct 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried going with this approach but was not able to get it to work.

I think this is technically non-blocking for now since we do not need the lambda ARN anymore for the CFN Primary Id.


static final AttributeKey<String> AWS_LAMBDA_RESOURCE_ID =
stringKey("aws.lambda.resource_mapping.id");

private AwsExperimentalAttributes() {}
}
26 changes: 26 additions & 0 deletions .../io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkExperimentalAttributesExtractor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,15 @@
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_ENDPOINT;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_GUARDRAIL_ID;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_KNOWLEDGE_BASE_ID;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_LAMBDA_NAME;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_LAMBDA_RESOURCE_ID;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_QUEUE_NAME;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_QUEUE_URL;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_REQUEST_ID;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_SECRET_ARN;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_SNS_TOPIC_ARN;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_STATE_MACHINE_ARN;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_STEP_FUNCTIONS_ACTIVITY_ARN;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_STREAM_NAME;
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_TABLE_NAME;

Expand Down Expand Up @@ -50,6 +56,18 @@ public void onStart(AttributesBuilder attributes, Context parentContext, Request
setAttribute(attributes, AWS_QUEUE_NAME, originalRequest, RequestAccess::getQueueName);
setAttribute(attributes, AWS_STREAM_NAME, originalRequest, RequestAccess::getStreamName);
setAttribute(attributes, AWS_TABLE_NAME, originalRequest, RequestAccess::getTableName);
setAttribute(
attributes, AWS_STATE_MACHINE_ARN, originalRequest, RequestAccess::getStateMachineArn);
setAttribute(
attributes,
AWS_STEP_FUNCTIONS_ACTIVITY_ARN,
originalRequest,
RequestAccess::getStepFunctionsActivityArn);
setAttribute(attributes, AWS_SNS_TOPIC_ARN, originalRequest, RequestAccess::getSnsTopicArn);
setAttribute(attributes, AWS_SECRET_ARN, originalRequest, RequestAccess::getSecretArn);
setAttribute(attributes, AWS_LAMBDA_NAME, originalRequest, RequestAccess::getLambdaName);
setAttribute(
attributes, AWS_LAMBDA_RESOURCE_ID, originalRequest, RequestAccess::getLambdaResourceId);

// Get serviceName defined in the AWS Java SDK V1 Request class.
String serviceName = request.getServiceName();
Expand All @@ -68,6 +86,14 @@ public void onEnd(
@Nullable Throwable error) {
if (response != null) {
Object awsResp = response.getAwsResponse();
setAttribute(attributes, AWS_STATE_MACHINE_ARN, awsResp, RequestAccess::getStateMachineArn);
setAttribute(
attributes,
AWS_STEP_FUNCTIONS_ACTIVITY_ARN,
awsResp,
RequestAccess::getStepFunctionsActivityArn);
setAttribute(attributes, AWS_SNS_TOPIC_ARN, awsResp, RequestAccess::getSnsTopicArn);
setAttribute(attributes, AWS_SECRET_ARN, awsResp, RequestAccess::getSecretArn);
if (awsResp instanceof AmazonWebServiceResponse) {
AmazonWebServiceResponse<?> awsWebServiceResponse = (AmazonWebServiceResponse<?>) awsResp;
String requestId = awsWebServiceResponse.getRequestId();
Expand Down
48 changes: 48 additions & 0 deletions ...11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/RequestAccess.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,42 @@ protected RequestAccess computeValue(Class<?> type) {
}
};

@Nullable
static String getLambdaName(Object request) {
RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
return invokeOrNull(access.getLambdaName, request);
}

@Nullable
static String getLambdaResourceId(Object request) {
RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
return invokeOrNull(access.getLambdaResourceId, request);
}

@Nullable
static String getSecretArn(Object request) {
RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
return invokeOrNull(access.getSecretArn, request);
}

@Nullable
static String getSnsTopicArn(Object request) {
RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
return invokeOrNull(access.getSnsTopicArn, request);
}

@Nullable
static String getStepFunctionsActivityArn(Object request) {
RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
return invokeOrNull(access.getStepFunctionsActivityArn, request);
}

@Nullable
static String getStateMachineArn(Object request) {
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());
Expand Down Expand Up @@ -102,6 +138,12 @@ private static String invokeOrNull(@Nullable MethodHandle method, Object obj) {
@Nullable private final MethodHandle getDataSourceId;
@Nullable private final MethodHandle getGuardrailId;
@Nullable private final MethodHandle getModelId;
@Nullable private final MethodHandle getStateMachineArn;
@Nullable private final MethodHandle getStepFunctionsActivityArn;
@Nullable private final MethodHandle getSnsTopicArn;
@Nullable private final MethodHandle getSecretArn;
@Nullable private final MethodHandle getLambdaName;
@Nullable private final MethodHandle getLambdaResourceId;

private RequestAccess(Class<?> clz) {
getBucketName = findAccessorOrNull(clz, "getBucketName");
Expand All @@ -114,6 +156,12 @@ private RequestAccess(Class<?> clz) {
getDataSourceId = findAccessorOrNull(clz, "getDataSourceId");
getGuardrailId = findAccessorOrNull(clz, "getGuardrailId");
getModelId = findAccessorOrNull(clz, "getModelId");
getStateMachineArn = findAccessorOrNull(clz, "getStateMachineArn");
getStepFunctionsActivityArn = findAccessorOrNull(clz, "getActivityArn");
getSnsTopicArn = findAccessorOrNull(clz, "getTopicArn");
getSecretArn = findAccessorOrNull(clz, "getARN");
getLambdaName = findAccessorOrNull(clz, "getFunctionName");
getLambdaResourceId = findAccessorOrNull(clz, "getUUID");
}

@Nullable
Expand Down
3 changes: 3 additions & 0 deletions instrumentation/aws-sdk/aws-sdk-1.11/testing/build.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ dependencies {
compileOnly("com.amazonaws:aws-java-sdk-bedrockagent:1.12.744")
compileOnly("com.amazonaws:aws-java-sdk-bedrockagentruntime:1.12.744")
compileOnly("com.amazonaws:aws-java-sdk-bedrockruntime:1.12.744")
compileOnly("com.amazonaws:aws-java-sdk-stepfunctions:1.11.+")
compileOnly("com.amazonaws:aws-java-sdk-secretsmanager:1.11.+")
compileOnly("com.amazonaws:aws-java-sdk-lambda:1.11.+")

// 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.12:1.0.0")
Expand Down
48 changes: 48 additions & 0 deletions ...c/main/groovy/io/opentelemetry/instrumentation/awssdk/v1_11/AbstractAws1ClientTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,15 @@ import com.amazonaws.services.bedrock.AmazonBedrockClientBuilder
import com.amazonaws.services.bedrock.model.GetGuardrailRequest
import com.amazonaws.services.bedrockruntime.AmazonBedrockRuntimeClientBuilder
import com.amazonaws.services.bedrockruntime.model.InvokeModelRequest
import com.amazonaws.services.stepfunctions.model.DescribeStateMachineRequest
import com.amazonaws.services.stepfunctions.model.DescribeActivityRequest
import com.amazonaws.services.sns.AmazonSNSClientBuilder
import com.amazonaws.services.sns.model.PublishRequest
import com.amazonaws.services.secretsmanager.AWSSecretsManagerClientBuilder
import com.amazonaws.services.secretsmanager.model.CreateSecretRequest
import com.amazonaws.services.lambda.AWSLambdaClientBuilder
import com.amazonaws.services.lambda.model.GetEventSourceMappingRequest
import com.amazonaws.services.lambda.model.GetFunctionRequest
import io.opentelemetry.api.trace.Span
import io.opentelemetry.instrumentation.test.InstrumentationSpecification
import io.opentelemetry.semconv.SemanticAttributes
Expand Down Expand Up @@ -207,6 +216,45 @@ abstract class AbstractAws1ClientTest extends InstrumentationSpecification {
"stop": "holes"
}
"""
"AWSStepFunctions" | "DescribeStateMachine" | "POST" | "/" | AWSStepFunctionsClientBuilder.standard()
| { c -> c.describeStateMachine(new DescribeStateMachineRequest().withStateMachineArn("stateMachineArn")) }
| ["aws.stepfunctions.state_machine.arn": "stateMachineArn"]
| ""
"AWSStepFunctions" | "DescribeActivity" | "POST" | "/" | AWSStepFunctionsClientBuilder.standard()
| { c -> c.describeActivity(new DescribeActivityRequest().withActivityArn("activityArn")) }
| ["aws.stepfunctions.activity.arn": "activityArn"]
| ""
"SNS" | "Publish" | "POST" | "/" | AmazonSNSClientBuilder.standard()
| { c -> c.publish(new PublishRequest().withMessage("message").withTopicArn("topicArn")) }
| ["aws.sns.topic.arn": "topicArn"]
| """
<PublishResponse xmlns="https://sns.amazonaws.com/doc/2010-03-31/">
<PublishResult>
<MessageId>567910cd-659e-55d4-8ccb-5aaf14679dc0</MessageId>
</PublishResult>
<ResponseMetadata>
<RequestId>d74b8436-ae13-5ab4-a9ff-ce54dfea72a0</RequestId>
</ResponseMetadata>
</PublishResponse>
"""
"AWSSecretsManager" | "CreateSecret" | "POST" | "/" | AWSSecretsManagerClientBuilder.standard()
| { c -> c.createSecret(new CreateSecretRequest().withName("secretName").withSecretString("secretValue")) }
| ["aws.secretsmanager.secret.arn": "arn:aws:secretsmanager:us-west-2:123456789012:secret:MyTestDatabaseSecret-a1b2c3"]
| """
{
"ARN": "arn:aws:secretsmanager:us-west-2:123456789012:secret:MyTestDatabaseSecret-a1b2c3",
"Name":"MyTestDatabaseSecret",
"VersionId": "EXAMPLE1-90ab-cdef-fedc-ba987SECRET1"
}
"""
"AWSLambda" | "GetEventSourceMapping" | "GET" | "/" | AWSLambdaClientBuilder.standard()
| { c -> c.getEventSourceMapping(new GetEventSourceMappingRequest().withUUID("uuid")) }
| ["aws.lambda.resource_mapping.id": "uuid"]
| ""
"AWSLambda" | "GetFunction" | "GET" | "/" | AWSLambdaClientBuilder.standard()
| { c-> c.getFunction(new GetFunctionRequest().withFunctionName("functionName")) }
| ["aws.lambda.function.name": "functionName"]
| ""
}

def "send #operation request to closed port"() {
Expand Down
Loading