Inject OpenTelemetrySdk into lambda library instrumentation instead o… (#2328)

* Inject OpenTelemetrySdk into lambda library instrumentation instead of using global.

* Update javadoc

* Drift

Author: Anuraag Agrawal

instrumentation/aws-lambda-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awslambda/v1_0/AwsLambdaInstrumentationHelper.java

@@ -5,18 +5,21 @@
 
 package io.opentelemetry.javaagent.instrumentation.awslambda.v1_0;
 
+import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.instrumentation.awslambda.v1_0.AwsLambdaMessageTracer;
 import io.opentelemetry.instrumentation.awslambda.v1_0.AwsLambdaTracer;
 
 public final class AwsLambdaInstrumentationHelper {
 
-  private static final AwsLambdaTracer FUNCTION_TRACER = new AwsLambdaTracer();
+  private static final AwsLambdaTracer FUNCTION_TRACER =
+      new AwsLambdaTracer(GlobalOpenTelemetry.get());
 
   public static AwsLambdaTracer functionTracer() {
     return FUNCTION_TRACER;
   }
 
-  private static final AwsLambdaMessageTracer MESSAGE_TRACER = new AwsLambdaMessageTracer();
+  private static final AwsLambdaMessageTracer MESSAGE_TRACER =
+      new AwsLambdaMessageTracer(GlobalOpenTelemetry.get());
 
   public static AwsLambdaMessageTracer messageTracer() {
     return MESSAGE_TRACER;

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

@@ -15,10 +15,20 @@ Available wrappers:
 
 ## Using handlers
 To use the instrumentation, replace your function classes that implement `RequestHandler` (or `RequestStreamHandler`) with those
-that extend `TracingRequestHandler` (or `TracingRequestStreamHandler`). You will need to change the method name to `doHandleRequest`.
+that extend `TracingRequestHandler` (or `TracingRequestStreamHandler`). You will need to change the method name to `doHandleRequest`
+and pass an initialized `OpenTelemetrySdk` to the base class.
 
 ```java
 public class MyRequestHandler extends TracingRequestHandler<String, String> {
+
+  private static final OpenTelemetrySdk SDK = OpenTelemetrySdk.builder()
+      .addSpanProcessor(spanProcessor)
+      .buildAndRegisterGlobal();
+
+  public MyRequestHandler() {
+    super(SDK);
+  }
+
   // Note the method is named doHandleRequest instead of handleRequest.
   @Override
   protected String doHandleRequest(String input, Context context) {
@@ -100,13 +110,15 @@ For API Gateway (HTTP) requests instrumented by using one of following methods:
 - wrapping with `TracingRequestStreamWrapper` or `TracingRequestApiGatewayWrapper`
 traces can be propagated with supported HTTP headers (see https://github.com/open-telemetry/opentelemetry-java/tree/master/extensions/trace_propagators).
 
-In order to enable requested propagation, configure it in your code as early as possible. For example B3 propagation configuration would look like as follows:
+In order to enable requested propagation for a handler, configure it on the SDK you build.
 
 ```java
   static {
-    OpenTelemetry.setGlobalPropagators(
-      DefaultContextPropagators.builder()
-        .addTextMapPropagator(B3Propagator.getInstance())
-        .build());
+    OpenTelemetrySdk.builder()
+      ...
+      .setPropagators(ContextPropagators.create(B3Propagator.injectingSingleHeader()))
+      .buildAndRegisterGlobal();
   }
 ```
+
+If using the wrappers, set the `OTEL_PROPAGATORS` environment variable as descibed [here](https://github.com/open-telemetry/opentelemetry-java/blob/main/sdk-extensions/autoconfigure/README.md#propagator).

instrumentation/aws-lambda-1.0/library/aws-lambda-1.0-library.gradle

@@ -17,6 +17,7 @@ apply from: "$rootDir/gradle/instrumentation-library.gradle"
 
 dependencies {
   compileOnly deps.opentelemetrySdk
+  compileOnly deps.opentelemetrySdkAutoconfigure
 
   library group: 'com.amazonaws', name: 'aws-lambda-java-core', version: '1.0.0'
   // First version to includes support for SQSEvent, currently the most popular message queue used
@@ -39,6 +40,7 @@ dependencies {
     'com.fasterxml.jackson.core:jackson-databind:2.10.5.1',
     'commons-io:commons-io:2.2')
 
+  testImplementation deps.opentelemetrySdkAutoconfigure
   testImplementation deps.opentelemetryTraceProps
   testImplementation deps.guava
 

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

@@ -7,11 +7,11 @@
 
 import com.amazonaws.services.lambda.runtime.events.SQSEvent;
 import com.amazonaws.services.lambda.runtime.events.SQSEvent.SQSMessage;
+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.Tracer;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
@@ -20,10 +20,8 @@ public class AwsLambdaMessageTracer extends BaseTracer {
 
   private static final String AWS_TRACE_HEADER_SQS_ATTRIBUTE_KEY = "AWSTraceHeader";
 
-  public AwsLambdaMessageTracer() {}
-
-  public AwsLambdaMessageTracer(Tracer tracer) {
-    super(tracer);
+  public AwsLambdaMessageTracer(OpenTelemetry openTelemetry) {
+    super(openTelemetry);
   }
 
   public Context startSpan(SQSEvent event) {

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

@@ -13,10 +13,10 @@
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
 import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent;
+import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanBuilder;
 import io.opentelemetry.api.trace.SpanKind;
-import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes.FaasTriggerValues;
 import java.lang.invoke.MethodHandle;
@@ -47,10 +47,8 @@ public class AwsLambdaTracer extends BaseTracer {
   // cached accountId value
   private volatile String accountId;
 
-  public AwsLambdaTracer() {}
-
-  public AwsLambdaTracer(Tracer tracer) {
-    super(tracer);
+  public AwsLambdaTracer(OpenTelemetry openTelemetry) {
+    super(openTelemetry);
   }
 
   private void setAttributes(SpanBuilder span, Context context, Object input) {
@@ -118,7 +116,7 @@ public io.opentelemetry.context.Context startSpan(
 
   public io.opentelemetry.context.Context startSpan(
       Context awsContext, SpanKind kind, Object input, Map<String, String> headers) {
-    io.opentelemetry.context.Context parentContext = ParentContextExtractor.extract(headers);
+    io.opentelemetry.context.Context parentContext = ParentContextExtractor.extract(headers, this);
 
     SpanBuilder spanBuilder = tracer.spanBuilder(spanName(awsContext, input));
     setAttributes(spanBuilder, awsContext, input);

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

@@ -20,15 +20,15 @@ public class ParentContextExtractor {
 
   private static final String AWS_TRACE_HEADER_ENV_KEY = "_X_AMZN_TRACE_ID";
 
-  static Context extract(Map<String, String> headers) {
+  static Context extract(Map<String, String> headers, BaseTracer tracer) {
     Context parentContext = null;
     String parentTraceHeader = System.getenv(AWS_TRACE_HEADER_ENV_KEY);
     if (parentTraceHeader != null) {
       parentContext = fromXRayHeader(parentTraceHeader);
     }
     if (!isValidAndSampled(parentContext)) {
       // try http
-      parentContext = fromHttpHeaders(headers);
+      parentContext = fromHttpHeaders(headers, tracer);
     }
     return parentContext;
   }
@@ -42,8 +42,8 @@ private static boolean isValidAndSampled(Context context) {
     return (parentSpanContext.isValid() && parentSpanContext.isSampled());
   }
 
-  private static Context fromHttpHeaders(Map<String, String> headers) {
-    return BaseTracer.extractWithGlobalPropagators(lowercaseMap(headers), MapGetter.INSTANCE);
+  private static Context fromHttpHeaders(Map<String, String> headers, BaseTracer tracer) {
+    return tracer.extract(lowercaseMap(headers), MapGetter.INSTANCE);
   }
 
   // lower-case map getter used for extraction

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

@@ -7,10 +7,21 @@
 
 import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
 import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
 
 /**
  * Wrapper for {@link TracingRequestHandler}. Allows for wrapping a lambda proxied through API
  * Gateway, enabling single span tracing and HTTP context propagation.
  */
 public class TracingRequestApiGatewayWrapper
-    extends TracingRequestWrapperBase<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {}
+    extends TracingRequestWrapperBase<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+  public TracingRequestApiGatewayWrapper() {
+    super();
+  }
+
+  // Visible for testing
+  TracingRequestApiGatewayWrapper(OpenTelemetrySdk openTelemetrySdk) {
+    super(openTelemetrySdk);
+  }
+}

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

@@ -9,10 +9,12 @@
 import com.amazonaws.services.lambda.runtime.RequestHandler;
 import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
 import io.opentelemetry.api.trace.SpanKind;
-import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import java.time.Duration;
 import java.util.Collections;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 /**
  * A base class similar to {@link RequestHandler} but will automatically trace invocations of {@link
@@ -21,38 +23,39 @@
  */
 public abstract class TracingRequestHandler<I, O> implements RequestHandler<I, O> {
 
-  private static final long DEFAULT_FLUSH_TIMEOUT_SECONDS = 1;
+  protected static final Duration DEFAULT_FLUSH_TIMEOUT = Duration.ofSeconds(1);
 
   private final AwsLambdaTracer tracer;
-  private final long flushTimeout;
+  private final OpenTelemetrySdk openTelemetrySdk;
+  private final long flushTimeoutNanos;
 
-  /** 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;
+  /**
+   * Creates a new {@link TracingRequestHandler} which traces using the provided {@link
+   * OpenTelemetrySdk} and has a timeout of 1s when flushing at the end of an invocation.
+   */
+  protected TracingRequestHandler(OpenTelemetrySdk openTelemetrySdk) {
+    this(openTelemetrySdk, DEFAULT_FLUSH_TIMEOUT);
   }
 
   /**
-   * Creates a new {@link TracingRequestHandler} which traces using the specified {@link Tracer}.
+   * Creates a new {@link TracingRequestHandler} which traces using the provided {@link
+   * OpenTelemetrySdk} and has a timeout of {@code flushTimeout} when flushing at the end of an
+   * invocation.
    */
-  protected TracingRequestHandler(Tracer tracer) {
-    this.tracer = new AwsLambdaTracer(tracer);
-    this.flushTimeout = DEFAULT_FLUSH_TIMEOUT_SECONDS;
+  protected TracingRequestHandler(OpenTelemetrySdk openTelemetrySdk, Duration flushTimeout) {
+    this(openTelemetrySdk, flushTimeout, new AwsLambdaTracer(openTelemetrySdk));
   }
 
   /**
-   * Creates a new {@link TracingRequestHandler} which traces using the specified {@link
-   * AwsLambdaTracer}.
+   * Creates a new {@link TracingRequestHandler} 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 AwsLambdaTracer}.
    */
-  protected TracingRequestHandler(AwsLambdaTracer tracer) {
+  protected TracingRequestHandler(
+      OpenTelemetrySdk openTelemetrySdk, Duration flushTimeout, AwsLambdaTracer tracer) {
+    this.openTelemetrySdk = openTelemetrySdk;
+    this.flushTimeoutNanos = flushTimeout.toNanos();
     this.tracer = tracer;
-    this.flushTimeout = DEFAULT_FLUSH_TIMEOUT_SECONDS;
   }
 
   private Map<String, String> getHeaders(I input) {
@@ -81,7 +84,10 @@ public final O handleRequest(I input, Context context) {
       } else {
         tracer.end(otelContext);
       }
-      LambdaUtils.forceFlush();
+      openTelemetrySdk
+          .getSdkTracerProvider()
+          .forceFlush()
+          .join(flushTimeoutNanos, TimeUnit.NANOSECONDS);
     }
   }