Let AWS Lambda SQS handlers report partial batch failures

Author: laurit

instrumentation/aws-lambda/aws-lambda-events-2.2/library/build.gradle.kts

@@ -8,9 +8,6 @@ dependencies {
   compileOnly("io.opentelemetry:opentelemetry-sdk")
   compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
 
-  compileOnly("com.google.auto.value:auto-value-annotations")
-  annotationProcessor("com.google.auto.value:auto-value")
-
   library("com.amazonaws:aws-lambda-java-core:1.0.0")
   // First version to includes support for SQSEvent, currently the most popular message queue used
   // with lambda.

instrumentation/aws-lambda/aws-lambda-events-2.2/library/src/main/java/io/opentelemetry/instrumentation/awslambdaevents/v2_2/TracingRequestWrapper.java

@@ -12,6 +12,7 @@
 import io.opentelemetry.instrumentation.awslambdacore.v1_0.internal.ApiGatewayProxyRequest;
 import io.opentelemetry.instrumentation.awslambdacore.v1_0.internal.MapUtils;
 import io.opentelemetry.instrumentation.awslambdacore.v1_0.internal.WrappedLambda;
+import io.opentelemetry.instrumentation.awslambdaevents.v2_2.internal.LambdaParameters;
 import io.opentelemetry.instrumentation.awslambdaevents.v2_2.internal.SerializationUtil;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import java.io.ByteArrayInputStream;

instrumentation/aws-lambda/aws-lambda-events-2.2/library/src/main/java/io/opentelemetry/instrumentation/awslambdaevents/v2_2/TracingRequestWrapperBase.java

@@ -13,6 +13,7 @@
 import io.opentelemetry.instrumentation.awslambdacore.v1_0.internal.WrappedLambda;
 import io.opentelemetry.instrumentation.awslambdacore.v1_0.internal.WrapperConfiguration;
 import io.opentelemetry.instrumentation.awslambdaevents.v2_2.internal.AwsLambdaEventsInstrumenterFactory;
+import io.opentelemetry.instrumentation.awslambdaevents.v2_2.internal.LambdaParameters;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
 import java.lang.reflect.InvocationTargetException;

instrumentation/aws-lambda/aws-lambda-events-2.2/library/src/main/java/io/opentelemetry/instrumentation/awslambdaevents/v2_2/TracingSqsEventWrapper.java

@@ -9,6 +9,7 @@
 import com.amazonaws.services.lambda.runtime.events.SQSEvent;
 import io.opentelemetry.instrumentation.awslambdacore.v1_0.internal.WrappedLambda;
 import io.opentelemetry.instrumentation.awslambdacore.v1_0.internal.WrapperConfiguration;
+import io.opentelemetry.instrumentation.awslambdaevents.v2_2.internal.LambdaParameters;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
 import java.lang.reflect.InvocationTargetException;

instrumentation/aws-lambda/aws-lambda-events-2.2/library/src/main/java/io/opentelemetry/instrumentation/awslambdaevents/v2_2/LambdaParameters.java renamed to instrumentation/aws-lambda/aws-lambda-events-2.2/library/src/main/java/io/opentelemetry/instrumentation/awslambdaevents/v2_2/internal/LambdaParameters.java

@@ -3,16 +3,20 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package io.opentelemetry.instrumentation.awslambdaevents.v2_2;
+package io.opentelemetry.instrumentation.awslambdaevents.v2_2.internal;
 
 import com.amazonaws.services.lambda.runtime.Context;
 import java.io.InputStream;
 import java.lang.reflect.Method;
 import java.util.function.BiFunction;
 
-final class LambdaParameters {
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public final class LambdaParameters {
 
-  static <T> Object[] toArray(
+  public static <T> Object[] toArray(
       Method targetMethod, T input, Context context, BiFunction<T, Class<?>, Object> mapper) {
     Class<?>[] parameterTypes = targetMethod.getParameterTypes();
     Object[] parameters = new Object[parameterTypes.length];
@@ -28,7 +32,7 @@ static <T> Object[] toArray(
     return parameters;
   }
 
-  static <T> Object[] toParameters(Method targetMethod, T input, Context context) {
+  public static <T> Object[] toParameters(Method targetMethod, T input, Context context) {
     Class<?>[] parameterTypes = targetMethod.getParameterTypes();
     Object[] parameters = new Object[parameterTypes.length];
     for (int i = 0; i < parameterTypes.length; i++) {
@@ -43,7 +47,7 @@ static <T> Object[] toParameters(Method targetMethod, T input, Context context)
     return parameters;
   }
 
-  static Object toInput(
+  public static Object toInput(
       Method targetMethod,
       InputStream inputStream,
       BiFunction<InputStream, Class<?>, Object> mapper) {

instrumentation/aws-lambda/aws-lambda-events-2.2/library/src/test/java/io/opentelemetry/instrumentation/awslambdaevents/v2_2/LambdaParametersTest.java renamed to instrumentation/aws-lambda/aws-lambda-events-2.2/library/src/test/java/io/opentelemetry/instrumentation/awslambdaevents/v2_2/internal/LambdaParametersTest.java

@@ -3,13 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package io.opentelemetry.instrumentation.awslambdaevents.v2_2;
+package io.opentelemetry.instrumentation.awslambdaevents.v2_2.internal;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
 
 import com.amazonaws.services.lambda.runtime.Context;
-import io.opentelemetry.instrumentation.awslambdaevents.v2_2.internal.SerializationUtil;
 import java.io.ByteArrayInputStream;
 import java.lang.reflect.Method;
 import org.junit.jupiter.api.Test;

instrumentation/aws-lambda/aws-lambda-events-2.2/testing/src/main/java/io/opentelemetry/instrumentation/awslambdaevents/v2_2/AbstractAwsLambdaSqsEventHandlerTest.java

@@ -35,7 +35,7 @@ public abstract class AbstractAwsLambdaSqsEventHandlerTest {
   private static final String AWS_TRACE_HEADER =
       "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1";
 
-  protected abstract RequestHandler<SQSEvent, Void> handler();
+  protected abstract RequestHandler<SQSEvent, ?> handler();
 
   protected abstract InstrumentationExtension testing();
 

instrumentation/aws-lambda/aws-lambda-events-3.11/library/build.gradle.kts

@@ -0,0 +1,30 @@
+plugins {
+  id("otel.library-instrumentation")
+}
+
+dependencies {
+  implementation(project(":instrumentation:aws-lambda:aws-lambda-events-2.2:library"))
+
+  compileOnly("io.opentelemetry:opentelemetry-sdk")
+  compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
+
+  library("com.amazonaws:aws-lambda-java-core:1.0.0")
+  library("com.amazonaws:aws-lambda-java-events:3.11.0")
+
+  // allows to get the function ARN
+  testLibrary("com.amazonaws:aws-lambda-java-core:1.2.1")
+
+  testImplementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
+  testImplementation("io.opentelemetry:opentelemetry-extension-trace-propagators")
+  testImplementation("com.amazonaws:aws-lambda-java-serialization:1.1.5")
+
+  testImplementation(project(":instrumentation:aws-lambda:aws-lambda-events-2.2:testing"))
+  testImplementation("uk.org.webcompere:system-stubs-jupiter")
+}
+
+tasks.withType<Test>().configureEach {
+  // required on jdk17
+  jvmArgs("--add-opens=java.base/java.lang=ALL-UNNAMED")
+  jvmArgs("--add-opens=java.base/java.util=ALL-UNNAMED")
+  jvmArgs("-XX:+IgnoreUnrecognizedVMOptions")
+}

instrumentation/aws-lambda/aws-lambda-events-3.11/library/src/main/java/io/opentelemetry/instrumentation/awslambdaevents/v3_11/TracingSqsEventHandler.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awslambdaevents.v3_11;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.events.SQSBatchResponse;
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.awslambdacore.v1_0.TracingRequestHandler;
+import io.opentelemetry.instrumentation.awslambdaevents.v2_2.internal.AwsLambdaSqsInstrumenterFactory;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import java.time.Duration;
+import javax.annotation.Nullable;
+
+public abstract class TracingSqsEventHandler
+    extends TracingRequestHandler<SQSEvent, SQSBatchResponse> {
+
+  private final Instrumenter<SQSEvent, Void> instrumenter;
+
+  /**
+   * Creates a new {@link TracingSqsEventHandler} which traces using the provided {@link
+   * OpenTelemetrySdk} and has a timeout of 1s when flushing at the end of an invocation.
+   */
+  protected TracingSqsEventHandler(OpenTelemetrySdk openTelemetrySdk) {
+    this(openTelemetrySdk, DEFAULT_FLUSH_TIMEOUT);
+  }
+
+  /**
+   * Creates a new {@link TracingSqsEventHandler} which traces using the provided {@link
+   * OpenTelemetrySdk} and has a timeout of {@code flushTimeout} when flushing at the end of an
+   * invocation.
+   */
+  protected TracingSqsEventHandler(OpenTelemetrySdk openTelemetrySdk, Duration flushTimeout) {
+    this(
+        openTelemetrySdk, flushTimeout, AwsLambdaSqsInstrumenterFactory.forEvent(openTelemetrySdk));
+  }
+
+  /**
+   * Creates a new {@link TracingSqsEventHandler} which flushes the provided {@link
+   * OpenTelemetrySdk}, has a timeout of {@code flushTimeout} when flushing at the end of an
+   * invocation, and traces using the provided {@link
+   * io.opentelemetry.instrumentation.awslambdacore.v1_0.internal.AwsLambdaFunctionInstrumenter}.
+   */
+  protected TracingSqsEventHandler(
+      OpenTelemetrySdk openTelemetrySdk,
+      Duration flushTimeout,
+      Instrumenter<SQSEvent, Void> instrumenter) {
+    super(openTelemetrySdk, flushTimeout);
+    this.instrumenter = instrumenter;
+  }
+
+  @Nullable
+  @Override
+  public SQSBatchResponse doHandleRequest(SQSEvent event, Context context) {
+    io.opentelemetry.context.Context parentContext = io.opentelemetry.context.Context.current();
+    if (instrumenter.shouldStart(parentContext, event)) {
+      io.opentelemetry.context.Context otelContext = instrumenter.start(parentContext, event);
+      Throwable error = null;
+      try (Scope ignored = otelContext.makeCurrent()) {
+        return handleEvent(event, context);
+      } catch (Throwable t) {
+        error = t;
+        throw t;
+      } finally {
+        instrumenter.end(otelContext, event, null, error);
+      }
+    } else {
+      return handleEvent(event, context);
+    }
+  }
+
+  /**
+   * Handles a {@linkplain SQSEvent batch of messages}. Implement this class to do the actual
+   * processing of incoming SQS messages.
+   */
+  @Nullable
+  protected abstract SQSBatchResponse handleEvent(SQSEvent event, Context context);
+}

instrumentation/aws-lambda/aws-lambda-events-3.11/library/src/main/java/io/opentelemetry/instrumentation/awslambdaevents/v3_11/TracingSqsEventWrapper.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awslambdaevents.v3_11;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.events.SQSBatchResponse;
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import io.opentelemetry.instrumentation.awslambdacore.v1_0.internal.WrappedLambda;
+import io.opentelemetry.instrumentation.awslambdacore.v1_0.internal.WrapperConfiguration;
+import io.opentelemetry.instrumentation.awslambdaevents.v2_2.internal.LambdaParameters;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+
+public class TracingSqsEventWrapper extends TracingSqsEventHandler {
+
+  private final WrappedLambda wrappedLambda;
+  private final Method targetMethod;
+
+  public TracingSqsEventWrapper() {
+    this(
+        AutoConfiguredOpenTelemetrySdk.initialize().getOpenTelemetrySdk(),
+        WrappedLambda.fromConfiguration());
+  }
+
+  // Visible for testing
+  TracingSqsEventWrapper(OpenTelemetrySdk openTelemetrySdk, WrappedLambda wrappedLambda) {
+    super(openTelemetrySdk, WrapperConfiguration.flushTimeout());
+    this.wrappedLambda = wrappedLambda;
+    this.targetMethod = wrappedLambda.getRequestTargetMethod();
+  }
+
+  @Override
+  protected SQSBatchResponse handleEvent(SQSEvent sqsEvent, Context context) {
+    Object[] parameters =
+        LambdaParameters.toArray(targetMethod, sqsEvent, context, (event, clazz) -> event);
+    try {
+      Object result = targetMethod.invoke(wrappedLambda.getTargetObject(), parameters);
+      return result instanceof SQSBatchResponse ? (SQSBatchResponse) result : null;
+    } catch (IllegalAccessException e) {
+      throw new IllegalStateException("Method is inaccessible", e);
+    } catch (InvocationTargetException e) {
+      throw (e.getCause() instanceof RuntimeException
+          ? (RuntimeException) e.getCause()
+          : new IllegalStateException(e.getTargetException()));
+    }
+  }
+}