Message handler

instrumentation/message-handler/library/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
+- `SQSBatchMessageHandlerTest` - processes messages from Amazon's SQS.
+
+## Using SQSBatchMessageHandlerTest
+
+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);
+```

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

@@ -0,0 +1,14 @@
+plugins {
+  id("otel.library-instrumentation")
+}
+
+dependencies {
+  api(project(":instrumentation:aws-lambda:aws-lambda-core-1.0:library"))
+
+  compileOnly("io.opentelemetry:opentelemetry-sdk")
+
+  library("com.amazonaws:aws-lambda-java-core:1.0.0")
+  library("com.amazonaws:aws-lambda-java-events:2.2.1")
+
+  implementation("io.opentelemetry.contrib:opentelemetry-aws-xray-propagator")
+}

instrumentation/message-handler/library/src/main/java/io/opentelemetry/instrumentation/messagehandler/BatchMessageHandler.java

@@ -0,0 +1,172 @@
+/*
+ * 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.SpanBuilder;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.TracerBuilder;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.contrib.awsxray.propagator.AwsXrayPropagator;
+import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessageOperation;
+import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Locale;
+import java.util.Map;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+public abstract class BatchMessageHandler<T> {
+  private static final Logger logger = Logger.getLogger(BatchMessageHandler.class.getName());
+  private static final String AWS_TRACE_HEADER_PROPAGATOR_KEY = "x-amzn-trace-id";
+
+  protected String messagingOperation;
+  protected OpenTelemetry openTelemetry;
+  protected String spanName;
+
+  public BatchMessageHandler(OpenTelemetry openTelemetry) {
+    this(openTelemetry, MessageOperation.RECEIVE.name());
+  }
+
+  public BatchMessageHandler(OpenTelemetry openTelemetry, String messageOperation) {
+    this(openTelemetry, messageOperation, "Batch Message");
+  }
+
+  public BatchMessageHandler(
+      OpenTelemetry openTelemetry, String messageOperation, String spanName) {
+    this.openTelemetry = openTelemetry;
+    this.spanName = spanName;
+    this.messagingOperation = messageOperation;
+  }
+
+  public abstract String getParentHeaderFromMessage(T t);
+
+  protected abstract void doHandleMessages(Collection<T> messages);
+
+  protected void addMessagingAttributes(SpanBuilder spanBuilder) {
+    spanBuilder.setAttribute(SemanticAttributes.MESSAGING_OPERATION, messagingOperation);
+  }
+
+  public void handleMessages(Collection<T> messages) {
+    TracerBuilder tracerBuilder =
+        openTelemetry
+            .tracerBuilder("io.opentelemetry.message.handler")
+            .setInstrumentationVersion("1.0");
+
+    Span parentSpan = Span.current();
+
+    SpanBuilder spanBuilder =
+        tracerBuilder
+            .build()
+            .spanBuilder(spanName)
+            .setParent(Context.current().with(parentSpan))
+            .setSpanKind(SpanKind.INTERNAL);
+
+    addMessagingAttributes(spanBuilder);
+
+    for (T t : messages) {
+      SpanContext spanContext = getParentSpanContextFromHeader(getParentHeaderFromMessage(t));
+
+      if (spanContext != null) {
+        spanBuilder.addLink(spanContext);
+      }
+    }
+
+    Span span = spanBuilder.startSpan();
+
+    try (Scope scope = span.makeCurrent()) {
+      doHandleMessages(messages);
+    } finally {
+      span.end();
+    }
+  }
+
+  public SpanContext getParentSpanContextFromHeader(String parentHeader) {
+    if (parentHeader == null) {
+      return null;
+    }
+
+    // We do not know if the upstream is W3C or X-Ray format.
+    // We will first try to decode it as a X-Ray trace context.
+    // Then we will try to decode it as a W3C trace context.
+    SpanContext spanContext = getParentSpanContextXray(parentHeader);
+
+    if (spanContext != null) {
+      return spanContext;
+    }
+
+    spanContext = getParentSpanContextW3C(parentHeader);
+
+    if (spanContext != null) {
+      return spanContext;
+    }
+
+    logger.log(Level.WARNING, "Invalid upstream span context: {0}", parentHeader);
+    return null;
+  }
+
+  private static SpanContext getParentSpanContextW3C(String parentHeader) {
+    try {
+      Context w3cContext =
+          W3CTraceContextPropagator.getInstance()
+              .extract(
+                  Context.root(),
+                  Collections.singletonMap("traceparent", parentHeader),
+                  MapGetter.INSTANCE);
+
+      SpanContext messageSpanCtx = Span.fromContext(w3cContext).getSpanContext();
+
+      if (messageSpanCtx.isValid()) {
+        return messageSpanCtx;
+      } else {
+        return null;
+      }
+    } catch (RuntimeException e) {
+      return null;
+    }
+  }
+
+  private static SpanContext getParentSpanContextXray(String parentHeader) {
+    try {
+      Context xrayContext =
+          AwsXrayPropagator.getInstance()
+              .extract(
+                  Context.root(),
+                  Collections.singletonMap(AWS_TRACE_HEADER_PROPAGATOR_KEY, parentHeader),
+                  MapGetter.INSTANCE);
+
+      SpanContext messageSpanCtx = Span.fromContext(xrayContext).getSpanContext();
+
+      if (messageSpanCtx.isValid()) {
+        return messageSpanCtx;
+      } else {
+        return null;
+      }
+    } catch (RuntimeException e) {
+      return null;
+    }
+  }
+
+  private enum MapGetter implements TextMapGetter<Map<String, String>> {
+    INSTANCE;
+
+    @Override
+    public Iterable<String> keys(Map<String, String> map) {
+      return map.keySet();
+    }
+
+    @Override
+    public String get(Map<String, String> map, String s) {
+      return map.get(s.toLowerCase(Locale.ROOT));
+    }
+  }
+}

instrumentation/message-handler/library/src/main/java/io/opentelemetry/instrumentation/messagehandler/SqsBatchMessageHandler.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.messagehandler;
+
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
+
+public abstract class SqsBatchMessageHandler extends BatchMessageHandler<SQSEvent.SQSMessage> {
+  private static final String AWS_TRACE_HEADER_SQS_ATTRIBUTE_KEY = "AWSTraceHeader";
+
+  public SqsBatchMessageHandler(OpenTelemetry openTelemetry) {
+    super(openTelemetry);
+  }
+
+  public SqsBatchMessageHandler(OpenTelemetry openTelemetry, String messageOperation) {
+    super(openTelemetry, messageOperation);
+  }
+
+  public SqsBatchMessageHandler(
+      OpenTelemetry openTelemetry, String messageOperation, String spanName) {
+    super(openTelemetry, messageOperation, spanName);
+  }
+
+  @Override
+  protected void addMessagingAttributes(SpanBuilder spanBuilder) {
+    super.addMessagingAttributes(spanBuilder);
+    spanBuilder.setAttribute(SemanticAttributes.MESSAGING_SYSTEM, "AmazonSQS");
+  }
+
+  @Override
+  public String getParentHeaderFromMessage(SQSEvent.SQSMessage message) {
+    return message.getAttributes().get(AWS_TRACE_HEADER_SQS_ATTRIBUTE_KEY);
+  }
+}