Skip to content

Commit

Permalink
Added dynamodb instrumenter for aws v1_11 sdk
Browse files Browse the repository at this point in the history
  • Loading branch information
Alex Kats authored and Alex Kats committed Nov 20, 2024
1 parent 40640a3 commit c711d74
Show file tree
Hide file tree
Showing 4 changed files with 69 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,7 @@ private <REQUEST, RESPONSE> AttributesExtractor<REQUEST, RESPONSE> messagingAttr
.build();
}


Instrumenter<SqsReceiveRequest, Response<?>> consumerReceiveInstrumenter() {
MessageOperation operation = MessageOperation.RECEIVE;
SqsReceiveRequestAttributesGetter getter = SqsReceiveRequestAttributesGetter.INSTANCE;
Expand Down Expand Up @@ -186,6 +187,18 @@ Instrumenter<Request<?>, Response<?>> producerInstrumenter() {
true);
}

Instrumenter<Request<?>, Response<?>> dynamoInstrumenter() {
DynamoAttributesExtractor dynamoAttributesExtractor = new DynamoAttributesExtractor();

return createInstrumenter(
openTelemetry,
spanName,
SpanKindExtractor.alwaysClient(),
attributesExtractors(),
singletonList(dynamoAttributesExtractor),
true);
}

private static <REQUEST, RESPONSE> Instrumenter<REQUEST, RESPONSE> createInstrumenter(
OpenTelemetry openTelemetry,
SpanNameExtractor<REQUEST> spanNameExtractor,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ public static AwsSdkTelemetryBuilder builder(OpenTelemetry openTelemetry) {
private final Instrumenter<SqsReceiveRequest, Response<?>> consumerReceiveInstrumenter;
private final Instrumenter<SqsProcessRequest, Response<?>> consumerProcessInstrumenter;
private final Instrumenter<Request<?>, Response<?>> producerInstrumenter;
private final Instrumenter<Request<?>, Response<?>> dynamoInstrumenter;

AwsSdkTelemetry(
OpenTelemetry openTelemetry,
Expand All @@ -65,6 +66,7 @@ public static AwsSdkTelemetryBuilder builder(OpenTelemetry openTelemetry) {
consumerReceiveInstrumenter = instrumenterFactory.consumerReceiveInstrumenter();
consumerProcessInstrumenter = instrumenterFactory.consumerProcessInstrumenter();
producerInstrumenter = instrumenterFactory.producerInstrumenter();
dynamoInstrumenter = instrumenterFactory.dynamoInstrumenter();
}

/**
Expand All @@ -76,6 +78,7 @@ public RequestHandler2 newRequestHandler() {
requestInstrumenter,
consumerReceiveInstrumenter,
consumerProcessInstrumenter,
producerInstrumenter);
producerInstrumenter,
dynamoInstrumenter);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
package io.opentelemetry.instrumentation.awssdk.v1_11;

import com.amazonaws.Request;
import com.amazonaws.Response;
import io.opentelemetry.api.common.AttributeKey;
import io.opentelemetry.api.common.AttributesBuilder;
import io.opentelemetry.context.Context;
import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
import io.opentelemetry.instrumentation.api.internal.AttributesExtractorUtil;
import javax.annotation.Nullable;
import java.util.Collections;
import java.util.List;

public class DynamoAttributesExtractor implements AttributesExtractor<Request<?>, Response<?>> {

private static final AttributeKey<String> DB_SYSTEM =
AttributeKey.stringKey("db.system");
private static final AttributeKey<List<String>> AWS_TABLE_NAMES =
AttributeKey.stringArrayKey("aws.dynamodb.table_names");

private static final String DYNAMODB = "dynamodb";

@Override
public void onStart(AttributesBuilder attributes, Context parentContext, Request<?> request) {
AttributesExtractorUtil.internalSet(attributes,DB_SYSTEM,DYNAMODB);
String tableName = RequestAccess.getTableName(request.getOriginalRequest());
AttributesExtractorUtil.internalSet(attributes,AWS_TABLE_NAMES, Collections.singletonList(tableName));
}

@Override
public void onEnd(
AttributesBuilder attributes,
Context context,
Request<?> request,
@Nullable Response<?> response,
@Nullable Throwable error) {}
}
Original file line number Diff line number Diff line change
Expand Up @@ -31,28 +31,34 @@ final class TracingRequestHandler extends RequestHandler2 {
ContextKey.named(TracingRequestHandler.class.getName() + ".Timer");
private static final ContextKey<Boolean> REQUEST_SPAN_SUPPRESSED_KEY =
ContextKey.named(TracingRequestHandler.class.getName() + ".RequestSpanSuppressed");
private static final String SEND_MESSAGE_REQUEST = "com.amazonaws.services.sqs.model.SendMessageRequest";
private static final String DYNAMODBV2 = "com.amazonaws.services.dynamodbv2.model";

private final Instrumenter<Request<?>, Response<?>> requestInstrumenter;
private final Instrumenter<SqsReceiveRequest, Response<?>> consumerReceiveInstrumenter;
private final Instrumenter<SqsProcessRequest, Response<?>> consumerProcessInstrumenter;
private final Instrumenter<Request<?>, Response<?>> producerInstrumenter;
private final Instrumenter<Request<?>, Response<?>> dynamoInstrumenter;

TracingRequestHandler(
Instrumenter<Request<?>, Response<?>> requestInstrumenter,
Instrumenter<SqsReceiveRequest, Response<?>> consumerReceiveInstrumenter,
Instrumenter<SqsProcessRequest, Response<?>> consumerProcessInstrumenter,
Instrumenter<Request<?>, Response<?>> producerInstrumenter) {
Instrumenter<Request<?>, Response<?>> producerInstrumenter,
Instrumenter<Request<?>, Response<?>> dynamoInstrumenter) {
this.requestInstrumenter = requestInstrumenter;
this.consumerReceiveInstrumenter = consumerReceiveInstrumenter;
this.consumerProcessInstrumenter = consumerProcessInstrumenter;
this.producerInstrumenter = producerInstrumenter;
this.dynamoInstrumenter = dynamoInstrumenter;
}

@Override
public void beforeRequest(Request<?> request) {
// GeneratePresignedUrlRequest doesn't result in actual request, beforeRequest is the only
// method called for it. Span created here would never be ended and scope would be leaked when
// running with java agent.
System.out.println(request.getOriginalRequest().getClass().getName());
if ("com.amazonaws.services.s3.model.GeneratePresignedUrlRequest"
.equals(request.getOriginalRequest().getClass().getName())) {
return;
Expand Down Expand Up @@ -151,14 +157,17 @@ private void finish(Request<?> request, Response<?> response, @Nullable Throwabl
}
return;
}

instrumenter.end(context, request, response, error);
}

private Instrumenter<Request<?>, Response<?>> getInstrumenter(Request<?> request) {
boolean isSqsProducer =
"com.amazonaws.services.sqs.model.SendMessageRequest"
.equals(request.getOriginalRequest().getClass().getName());
return isSqsProducer ? producerInstrumenter : requestInstrumenter;
String className = request.getOriginalRequest().getClass().getName();
if (className.startsWith(DYNAMODBV2)) {
return dynamoInstrumenter;
}
else if (className.equals(SEND_MESSAGE_REQUEST)) {
return producerInstrumenter;
}
return requestInstrumenter;
}
}

0 comments on commit c711d74

Please sign in to comment.