AWS lambda - configurable flush timeout (#1960)

instrumentation/aws-lambda-1.0/library/README.md

@@ -3,9 +3,11 @@
 This package contains libraries to help instrument AWS lambda functions in your code.
 
 ## Using wrappers
-To use the instrumentation, configure `OTEL_LAMBDA_HANDLER` env property to your lambda handler method in following format `package.ClassName::methodName`
+To use the instrumentation, configure `OTEL_INSTRUMENTATION_AWS_LAMBDA_HANDLER` env property to your lambda handler method in following format `package.ClassName::methodName`
 and use one of wrappers as your lambda `Handler`.
 
+In order to configure a span flush timeout (default is set to 1 second), please configure `OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT` env property. The value is in seconds.
+
 Available wrappers:
 - `io.opentelemetry.instrumentation.awslambda.v1_0.TracingRequestWrapper` - for wrapping regular handlers (implementing `RequestHandler`)
 - `io.opentelemetry.instrumentation.awslambda.v1_0.TracingRequestApiGatewayWrapper` - for wrapping regular handlers (implementing `RequestHandler`) proxied through API Gateway, enabling HTTP context propagation

instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingRequestHandler.java

@@ -24,11 +24,21 @@ import java.util.concurrent.TimeUnit;
  */
 public abstract class TracingRequestHandler<I, O> implements RequestHandler<I, O> {
 
+  private static final long DEFAULT_FLUSH_TIMEOUT_SECONDS = 1;
+
   private final AwsLambdaTracer tracer;
+  private final long flushTimeout;
+
+  /** Creates a new {@link TracingRequestHandler} which traces using the default {@link Tracer}. */
+  protected TracingRequestHandler(long flushTimeout) {
+    this.tracer = new AwsLambdaTracer();
+    this.flushTimeout = flushTimeout;
+  }
 
   /** Creates a new {@link TracingRequestHandler} which traces using the default {@link Tracer}. */
   protected TracingRequestHandler() {
     this.tracer = new AwsLambdaTracer();
+    this.flushTimeout = DEFAULT_FLUSH_TIMEOUT_SECONDS;
   }
 
   /**
@@ -36,6 +46,7 @@ public abstract class TracingRequestHandler<I, O> implements RequestHandler<I, O
    */
   protected TracingRequestHandler(Tracer tracer) {
     this.tracer = new AwsLambdaTracer(tracer);
+    this.flushTimeout = DEFAULT_FLUSH_TIMEOUT_SECONDS;
   }
 
   /**
@@ -44,6 +55,7 @@ public abstract class TracingRequestHandler<I, O> implements RequestHandler<I, O
    */
   protected TracingRequestHandler(AwsLambdaTracer tracer) {
     this.tracer = tracer;
+    this.flushTimeout = DEFAULT_FLUSH_TIMEOUT_SECONDS;
   }
 
   private Map<String, String> getHeaders(I input) {
@@ -71,7 +83,9 @@ public abstract class TracingRequestHandler<I, O> implements RequestHandler<I, O
       } else {
         tracer.end(span);
       }
-      OpenTelemetrySdk.getGlobalTracerManagement().forceFlush().join(1, TimeUnit.SECONDS);
+      OpenTelemetrySdk.getGlobalTracerManagement()
+          .forceFlush()
+          .join(flushTimeout, TimeUnit.SECONDS);
     }
   }
 

instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingRequestStreamHandler.java

@@ -23,7 +23,19 @@ import java.util.concurrent.TimeUnit;
  */
 public abstract class TracingRequestStreamHandler implements RequestStreamHandler {
 
+  private static final long DEFAULT_FLUSH_TIMEOUT_SECONDS = 1;
+
   private final AwsLambdaTracer tracer;
+  private final long flushTimeout;
+
+  /**
+   * Creates a new {@link TracingRequestStreamHandler} which traces using the default {@link
+   * Tracer}.
+   */
+  protected TracingRequestStreamHandler(long flushTimeout) {
+    this.tracer = new AwsLambdaTracer();
+    this.flushTimeout = flushTimeout;
+  }
 
   /**
    * Creates a new {@link TracingRequestStreamHandler} which traces using the default {@link
@@ -31,6 +43,7 @@ public abstract class TracingRequestStreamHandler implements RequestStreamHandle
    */
   protected TracingRequestStreamHandler() {
     this.tracer = new AwsLambdaTracer();
+    this.flushTimeout = DEFAULT_FLUSH_TIMEOUT_SECONDS;
   }
 
   /**
@@ -39,6 +52,7 @@ public abstract class TracingRequestStreamHandler implements RequestStreamHandle
    */
   protected TracingRequestStreamHandler(Tracer tracer) {
     this.tracer = new AwsLambdaTracer(tracer);
+    this.flushTimeout = DEFAULT_FLUSH_TIMEOUT_SECONDS;
   }
 
   /**
@@ -47,6 +61,7 @@ public abstract class TracingRequestStreamHandler implements RequestStreamHandle
    */
   protected TracingRequestStreamHandler(AwsLambdaTracer tracer) {
     this.tracer = tracer;
+    this.flushTimeout = DEFAULT_FLUSH_TIMEOUT_SECONDS;
   }
 
   @Override
@@ -60,7 +75,9 @@ public abstract class TracingRequestStreamHandler implements RequestStreamHandle
       doHandleRequest(proxyRequest.freshStream(), new OutputStreamWrapper(output, span), context);
     } catch (Throwable t) {
       tracer.endExceptionally(span, t);
-      OpenTelemetrySdk.getGlobalTracerManagement().forceFlush().join(1, TimeUnit.SECONDS);
+      OpenTelemetrySdk.getGlobalTracerManagement()
+          .forceFlush()
+          .join(flushTimeout, TimeUnit.SECONDS);
       throw t;
     }
   }

instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingRequestStreamWrapper.java

@@ -13,15 +13,19 @@ import java.io.OutputStream;
 
 /**
  * Wrapper for {@link TracingRequestStreamHandler}. Allows for wrapping a regular lambda, enabling
- * single span tracing. Main lambda class should be configured as env property OTEL_LAMBDA_HANDLER
+ * single span tracing. Main lambda class should be configured as env property
- * in package.ClassName::methodName format. Lambda class must implement {@link
+ * OTEL_INSTRUMENTATION_AWS_LAMBDA_HANDLER in package.ClassName::methodName format. Lambda class
- * RequestStreamHandler}.
+ * must implement {@link RequestStreamHandler}.
  */
 public class TracingRequestStreamWrapper extends TracingRequestStreamHandler {
 
   // visible for testing
   static WrappedLambda WRAPPED_LAMBDA = WrappedLambda.fromConfiguration();
 
+  public TracingRequestStreamWrapper() {
+    super(WrapperConfiguration.flushTimeout());
+  }
+
   @Override
   protected void doHandleRequest(InputStream inputStream, OutputStream output, Context context)
       throws IOException {

instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingRequestWrapperBase.java

@@ -11,10 +11,14 @@ import java.lang.reflect.Method;
 
 /**
  * Base abstract wrapper for {@link TracingRequestHandler}. Provides: - delegation to a lambda via
- * env property OTEL_LAMBDA_HANDLER in package.ClassName::methodName format
+ * env property OTEL_INSTRUMENTATION_AWS_LAMBDA_HANDLER in package.ClassName::methodName format
  */
 abstract class TracingRequestWrapperBase<I, O> extends TracingRequestHandler<I, O> {
 
+  protected TracingRequestWrapperBase() {
+    super(WrapperConfiguration.flushTimeout());
+  }
+
   // visible for testing
   static WrappedLambda WRAPPED_LAMBDA = WrappedLambda.fromConfiguration();
 

instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/WrappedLambda.java

@@ -17,7 +17,8 @@ import java.util.Optional;
 /** Model for wrapped lambda function (object, class, method). */
 class WrappedLambda {
 
-  public static final String OTEL_LAMBDA_HANDLER_ENV_KEY = "OTEL_LAMBDA_HANDLER";
+  public static final String OTEL_LAMBDA_HANDLER_ENV_KEY =
+      "OTEL_INSTRUMENTATION_AWS_LAMBDA_HANDLER";
 
   private final Object targetObject;
   private final Class<?> targetClass;

instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/WrapperConfiguration.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awslambda.v1_0;
+
+public final class WrapperConfiguration {
+
+  private WrapperConfiguration() {}
+
+  public static final String OTEL_LAMBDA_FLUSH_TIMEOUT_ENV_KEY =
+      "OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT";
+  public static final long OTEL_LAMBDA_FLUSH_TIMEOUT_DEFAULT = 1;
+
+  public static final long flushTimeout() {
+    String lambdaFlushTimeout = System.getenv(OTEL_LAMBDA_FLUSH_TIMEOUT_ENV_KEY);
+    if (lambdaFlushTimeout != null && !lambdaFlushTimeout.isEmpty()) {
+      try {
+        return Long.parseLong(lambdaFlushTimeout);
+      } catch (NumberFormatException nfe) {
+        // ignored - default used
+      }
+    }
+    return OTEL_LAMBDA_FLUSH_TIMEOUT_DEFAULT;
+  }
+}