Added dynamodb instrumenter for aws v1_11 sdk

open-telemetry · Nov 20, 2024 · fceec97 · fceec97
1 parent 40640a3
commit fceec97
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 7 deletions.
diff --git a/...rc/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkInstrumenterFactory.java b/...rc/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkInstrumenterFactory.java
@@ -103,6 +103,7 @@ private <REQUEST, RESPONSE> AttributesExtractor<REQUEST, RESPONSE> messagingAttr
         .build();
   }
 
+
   Instrumenter<SqsReceiveRequest, Response<?>> consumerReceiveInstrumenter() {
     MessageOperation operation = MessageOperation.RECEIVE;
     SqsReceiveRequestAttributesGetter getter = SqsReceiveRequestAttributesGetter.INSTANCE;
@@ -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,

diff --git a/.../library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkTelemetry.java b/.../library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkTelemetry.java
@@ -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,
@@ -65,6 +66,7 @@ public static AwsSdkTelemetryBuilder builder(OpenTelemetry openTelemetry) {
     consumerReceiveInstrumenter = instrumenterFactory.consumerReceiveInstrumenter();
     consumerProcessInstrumenter = instrumenterFactory.consumerProcessInstrumenter();
     producerInstrumenter = instrumenterFactory.producerInstrumenter();
+    dynamoInstrumenter = instrumenterFactory.dynamoInstrumenter();
   }
 
   /**
@@ -76,6 +78,7 @@ public RequestHandler2 newRequestHandler() {
         requestInstrumenter,
         consumerReceiveInstrumenter,
         consumerProcessInstrumenter,
-        producerInstrumenter);
+        producerInstrumenter,
+        dynamoInstrumenter);
   }
 }
diff --git a/...rc/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/DynamoAttributesExtractor.java b/...rc/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/DynamoAttributesExtractor.java
@@ -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) {}
+}
diff --git a/...ry/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java b/...ry/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java
@@ -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;
@@ -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;
   }
 }