Message handler

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/SqsMessageAccess.java

@@ -30,7 +30,7 @@
  *     href="https://github.com/aws/aws-sdk-java-v2/blob/2.2.0/services/sqs/src/main/resources/codegen-resources/service-2.json#L821-L856">Definition
  *     JSON</a>
  */
-final class SqsMessageAccess {
+public final class SqsMessageAccess {
 
   @Nullable private static final MethodHandle GET_ATTRIBUTES;
   @Nullable private static final MethodHandle GET_MESSAGE_ATTRIBUTES;
@@ -68,7 +68,7 @@ final class SqsMessageAccess {
   }
 
   @SuppressWarnings("unchecked")
-  static Map<String, String> getAttributes(Object message) {
+  public static Map<String, String> getAttributes(Object message) {
     if (GET_ATTRIBUTES == null) {
       return Collections.emptyMap();
     }

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/SqsParentContext.java

@@ -13,7 +13,7 @@
 import java.util.Map;
 import software.amazon.awssdk.core.SdkPojo;
 
-final class SqsParentContext {
+public final class SqsParentContext {
 
   enum StringMapGetter implements TextMapGetter<Map<String, String>> {
     INSTANCE;
@@ -52,13 +52,13 @@ public String get(Map<String, SdkPojo> map, String s) {
 
   static final String AWS_TRACE_SYSTEM_ATTRIBUTE = "AWSTraceHeader";
 
-  static Context ofMessageAttributes(
+  public static Context ofMessageAttributes(
       Map<String, SdkPojo> messageAttributes, TextMapPropagator propagator) {
     return propagator.extract(
         Context.root(), messageAttributes, MessageAttributeValueMapGetter.INSTANCE);
   }
 
-  static Context ofSystemAttributes(Map<String, String> systemAttributes) {
+  public static Context ofSystemAttributes(Map<String, String> systemAttributes) {
     String traceHeader = systemAttributes.get(AWS_TRACE_SYSTEM_ATTRIBUTE);
     return AwsXrayPropagator.getInstance()
         .extract(

instrumentation/message-handler/aws/README.md

@@ -0,0 +1,19 @@
+# Using SqsMessageHandler
+
+1. Retrieve a collection of messages to process.
+2. Create a SqsMessageHandler and provide the business logic on what to do with the messages.
+3. Call the handleMessages function and pass in your messages.
+4. It will call the doHandleMessages function you provided wrapped in the messaging span.
+
+```java
+Collection<SQSEvent.SQSMessage> sqsMessages;
+
+SqsMessageHandler messageHandler = new SqsMessageHandler(opentelemetry) {
+  @Override
+  protected void doHandleMessages(Collection<SQSEvent.SQSMessage> messages) {
+    // Do my business logic
+  }
+};
+
+messageHandler.handleMessages(sqsMessages, "destination.name", MessageOperation.RECEIVE);
+```

instrumentation/message-handler/aws/build.gradle.kts

@@ -0,0 +1,14 @@
+plugins {
+  id("otel.library-instrumentation")
+}
+
+dependencies {
+  compileOnly("io.opentelemetry:opentelemetry-sdk")
+
+  implementation(project(":instrumentation:message-handler:core"))
+  implementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:library"))
+
+  library("software.amazon.awssdk:sqs:2.2.0")
+
+  implementation("io.opentelemetry.contrib:opentelemetry-aws-xray-propagator")
+}

instrumentation/message-handler/aws/src/main/java/io/opentelemetry/instrumentation/messagehandler/SqsMessageHandler.java

@@ -0,0 +1,160 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.messagehandler;
+
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanLinksExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessageOperation;
+import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessagingAttributesExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessagingAttributesGetter;
+import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessagingSpanNameExtractor;
+import io.opentelemetry.instrumentation.awssdk.v2_2.SqsMessageAccess;
+import io.opentelemetry.instrumentation.awssdk.v2_2.SqsParentContext;
+import java.util.Collection;
+import java.util.Map;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.Nullable;
+import software.amazon.awssdk.core.SdkPojo;
+import software.amazon.awssdk.services.sqs.model.Message;
+
+public abstract class SqsMessageHandler extends MessageHandler<Message> {
+  private static final Logger logger = Logger.getLogger(SqsMessageHandler.class.getName());
+
+  private final OpenTelemetry openTelemetry;
+  private final MessageOperation messageOperation;
+  private final String destination;
+
+  public SqsMessageHandler(
+      OpenTelemetry openTelemetry, String destination, MessageOperation messageOperation) {
+    this.openTelemetry = openTelemetry;
+    this.messageOperation = messageOperation;
+    this.destination = destination;
+  }
+
+  @Override
+  protected Instrumenter<Collection<Message>, Void> getMessageInstrumenter() {
+    return Instrumenter.<Collection<Message>, Void>builder(
+            openTelemetry, "io.opentelemetry.message-handler", getSpanNameExtractor())
+        .setInstrumentationVersion("1.0")
+        .addAttributesExtractor(getAttributesExtractor())
+        .addSpanLinksExtractor(getSpanLinksExtractor())
+        .buildInstrumenter(getSpanKindExtractor());
+  }
+
+  protected SpanNameExtractor<Collection<Message>> getSpanNameExtractor() {
+    return MessagingSpanNameExtractor.create(getMessageingAttributesGetter(), messageOperation);
+  }
+
+  private MessagingAttributesGetter<Collection<Message>, Void> getMessageingAttributesGetter() {
+    String destination = this.destination;
+
+    return new MessagingAttributesGetter<Collection<Message>, Void>() {
+      @Nullable
+      @Override
+      public String getSystem(Collection<Message> messages) {
+        return "AmazonSQS";
+      }
+
+      @Nullable
+      @Override
+      public String getDestinationKind(Collection<Message> messages) {
+        return null;
+      }
+
+      @Nullable
+      @Override
+      public String getDestination(Collection<Message> messages) {
+        return destination;
+      }
+
+      @Override
+      public boolean isTemporaryDestination(Collection<Message> messages) {
+        return false;
+      }
+
+      @Nullable
+      @Override
+      public String getConversationId(Collection<Message> messages) {
+        return null;
+      }
+
+      @Nullable
+      @Override
+      public Long getMessagePayloadSize(Collection<Message> messages) {
+        long total = 0;
+
+        for (Message message : messages) {
+          if (message.body() != null) {
+            total += message.body().length();
+          }
+        }
+
+        return total;
+      }
+
+      @Nullable
+      @Override
+      public Long getMessagePayloadCompressedSize(Collection<Message> messages) {
+        return null;
+      }
+
+      @Nullable
+      @Override
+      public String getMessageId(Collection<Message> messages, @Nullable Void unused) {
+        return null;
+      }
+    };
+  }
+
+  protected SpanKindExtractor<Collection<Message>> getSpanKindExtractor() {
+    if (messageOperation == MessageOperation.RECEIVE
+        || messageOperation == MessageOperation.PROCESS) {
+      return SpanKindExtractor.alwaysConsumer();
+    } else if (messageOperation == MessageOperation.SEND) {
+      return SpanKindExtractor.alwaysProducer();
+    } else {
+      logger.log(
+          Level.WARNING, "Unknown Messaging Operation {0}", new Object[] {messageOperation.name()});
+      return SpanKindExtractor.alwaysConsumer();
+    }
+  }
+
+  protected AttributesExtractor<Collection<Message>, Void> getAttributesExtractor() {
+    return MessagingAttributesExtractor.create(getMessageingAttributesGetter(), messageOperation);
+  }
+
+  protected SpanLinksExtractor<Collection<Message>> getSpanLinksExtractor() {
+    TextMapPropagator messagingPropagator = openTelemetry.getPropagators().getTextMapPropagator();
+
+    return (spanLinks, parentContext, sqsMessages) -> {
+      for (Message message : sqsMessages) {
+        Map<String, SdkPojo> messageAtributes = SqsMessageAccess.getMessageAttributes(message);
+
+        Context context =
+            SqsParentContext.ofMessageAttributes(messageAtributes, messagingPropagator);
+
+        if (context == Context.root()) {
+          context = SqsParentContext.ofSystemAttributes(SqsMessageAccess.getAttributes(message));
+        }
+
+        SpanContext messageSpanCtx = Span.fromContext(context).getSpanContext();
+
+        if (messageSpanCtx.isValid()) {
+          spanLinks.addLink(messageSpanCtx);
+        }
+      }
+    };
+  }
+}