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,30 @@
+# Message Handler
+
+This package contains handlers to instrument message system spans for a batch of messages. It is not designed for a for-loop case.
+
+The handler will create a single new messaging span, add span links to it, and set the messaging attributes behind the scene. These values are based off of the messages passed in.
+
+The handler provides constructors to change the messaging operation and span name of the newly created messaging span.
+
+## Available Handlers
+- `SQSBatchMessageHandler` - processes messages from Amazon's SQS.
+
+## Using SQSBatchMessageHandler
+
+1. Retrieve a collection of messages to process.
+2. Create a SQSBatchMessageHandler 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;
+
+SqsBatchMessageHandler messageHandler = new SqsBatchMessageHandler(GlobalOpenTelemetry.get()) {
+  @Override
+  protected void doHandleMessages(Collection<SQSEvent.SQSMessage> messages) {
+    // Do my business logic
+  }
+};
+
+messageHandler.handleMessages(sqsMessages, messages -> "Receiving Messages");
+```

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/SqsBatchMessageHandler.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.messagehandler;
+
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.common.AttributesBuilder;
+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.SpanLinksExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import io.opentelemetry.instrumentation.awssdk.v2_2.SqsMessageAccess;
+import io.opentelemetry.instrumentation.awssdk.v2_2.SqsParentContext;
+import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
+import java.util.Collection;
+import java.util.Map;
+import javax.annotation.Nullable;
+import software.amazon.awssdk.core.SdkPojo;
+import software.amazon.awssdk.services.sqs.model.Message;
+
+public abstract class SqsBatchMessageHandler extends BatchMessageHandler<Message> {
+  public SqsBatchMessageHandler(
+      OpenTelemetry openTelemetry, SpanNameExtractor<Collection<Message>> spanNameExtractor) {
+    super(openTelemetry, spanNameExtractor);
+  }
+
+  public SqsBatchMessageHandler(
+      OpenTelemetry openTelemetry,
+      SpanNameExtractor<Collection<Message>> spanNameExtractor,
+      String messageOperation) {
+    super(openTelemetry, spanNameExtractor, messageOperation);
+  }
+
+  @Override
+  protected void setup() {
+    messageInstrumenter =
+        Instrumenter.<Collection<Message>, Void>builder(
+                getOpenTelemetry(), "io.opentelemetry.message.handler", getSpanNameExtractor())
+            .setInstrumentationVersion("1.0")
+            .addAttributesExtractor(getGenericAttributesExtractor())
+            .addAttributesExtractor(getAttributesExtractor())
+            .addSpanLinksExtractor(getSpanLinksExtractor())
+            .buildInstrumenter();
+  }
+
+  protected AttributesExtractor<Collection<Message>, Void> getAttributesExtractor() {
+    return new AttributesExtractor<Collection<Message>, Void>() {
+
+      @Override
+      public void onStart(
+          AttributesBuilder attributes, Context parentContext, Collection<Message> messages) {
+        attributes.put(SemanticAttributes.MESSAGING_SYSTEM, "AmazonSQS");
+      }
+
+      @Override
+      public void onEnd(
+          AttributesBuilder attributes,
+          Context context,
+          Collection<Message> messages,
+          @Nullable Void unused,
+          @Nullable Throwable error) {}
+    };
+  }
+
+  protected SpanLinksExtractor<Collection<Message>> getSpanLinksExtractor() {
+    TextMapPropagator messagingPropagator =
+        getOpenTelemetry().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);
+        }
+      }
+    };
+  }
+}