Refactor TracingOpenTelemetry for handler-specific attributes, response/error capture, and contextual propagation

Author: David-DAM

powertools-tracing-opentelemetry/pom.xml

@@ -28,7 +28,8 @@
 
     <name>Powertools for AWS Lambda (Java) - Tracing OpenTelemetry</name>
     <description>
-        A suite of utilities for AWS Lambda Functions that makes tracing with OpenTelemetry, structured logging and creating custom metrics asynchronously easier.
+        A suite of utilities for AWS Lambda Functions that makes tracing with OpenTelemetry, structured logging and
+        creating custom metrics asynchronously easier.
     </description>
 
     <dependencies>
@@ -37,6 +38,16 @@
             <artifactId>opentelemetry-api</artifactId>
             <version>${opentelemetry-api.version}</version>
         </dependency>
+        <dependency>
+            <groupId>io.opentelemetry</groupId>
+            <artifactId>opentelemetry-sdk</artifactId>
+            <version>${opentelemetry-api.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>io.opentelemetry</groupId>
+            <artifactId>opentelemetry-exporter-otlp</artifactId>
+            <version>${opentelemetry-api.version}</version>
+        </dependency>
         <dependency>
             <groupId>org.aspectj</groupId>
             <artifactId>aspectjrt</artifactId>
@@ -64,12 +75,6 @@
         </dependency>
 
         <!-- Test dependencies -->
-        <dependency>
-            <groupId>io.opentelemetry</groupId>
-            <artifactId>opentelemetry-sdk</artifactId>
-            <version>${opentelemetry-api.version}</version>
-            <scope>test</scope>
-        </dependency>
         <dependency>
             <groupId>io.opentelemetry</groupId>
             <artifactId>opentelemetry-sdk-testing</artifactId>

…s/tracing/opentelemetry/TracingOtel.java …tools/tracing/opentelemetry/Tracing.javapowertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/TracingOtel.java renamed to powertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/Tracing.java powertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/TracingOtel.java renamed to powertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/Tracing.java

@@ -14,7 +14,7 @@
  */
 @Retention(RetentionPolicy.RUNTIME)
 @Target(ElementType.METHOD)
-public @interface TracingOtel {
+public @interface Tracing {
     /**
      * The namespace associated with the span.
      *

powertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/TracingOpenTelemetry.java

@@ -2,8 +2,10 @@
  * Copyright 2023 Amazon.com, Inc. or its affiliates.
  * Licensed under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
+ * with the License. You may obtain a copy of the License at
+ *
  *     http://www.apache.org/licenses/LICENSE-2.0
+ *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -14,100 +16,240 @@
 
 package software.amazon.lambda.powertools.tracing.opentelemetry;
 
-import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.context.propagation.TextMapSetter;
 import java.util.Objects;
+import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
+import software.amazon.lambda.powertools.tracing.opentelemetry.internal.OpenTelemetryProvider;
 import software.amazon.lambda.powertools.tracing.opentelemetry.internal.SpanOperation;
 import software.amazon.lambda.powertools.tracing.opentelemetry.internal.SpanScope;
 
-/**
- * A wrapper for OpenTelemetry's Tracer that simplifies the creation and management of spans.
- * This class provides utility functions for starting and controlling spans and their contexts
- * in the current execution thread.
- * <p>
- * This is a final class and cannot be extended.
- */
-public final class TracingOpenTelemetry {
 
-    private static final String INSTRUMENTATION_NAME =
-            "aws-lambda-powertools";
+public final class TracingOpenTelemetry {
 
     private final Tracer tracer;
+    private final TextMapPropagator propagator;
+
+    private TracingOpenTelemetry(Builder builder) {
+        this.tracer = Objects.requireNonNull(
+                builder.tracer,
+                "tracer must not be null"
+        );
+        this.propagator = Objects.requireNonNull(
+                builder.propagator,
+                "propagator must not be null"
+        );
+    }
 
-    /**
-     * Creates a tracing instance using the provided tracer.
-     *
-     * <p>This constructor is primarily useful for testing.
-     *
-     * @param tracer the OpenTelemetry tracer
-     */
-    TracingOpenTelemetry(Tracer tracer) {
-        this.tracer = Objects.requireNonNull(tracer, "tracer must not be null");
-    }
-
-    /**
-     * Initializes a new instance of the {@code TracingOpenTelemetry} class, using
-     * the global OpenTelemetry tracer identified by the instrumentation name.
-     * <p>
-     * This constructor simplifies the setup process for applications by
-     * automatically leveraging the globally configured instrumentation tracer.
-     */
     public TracingOpenTelemetry() {
-        this(GlobalOpenTelemetry.getTracer(INSTRUMENTATION_NAME));
+        this(OpenTelemetryProvider.tracer());
+    }
+
+
+    public TracingOpenTelemetry(Tracer tracer) {
+        this(tracer, createDefaultPropagator());
     }
 
-    /**
-     * Creates a new span with the specified name and makes it the current span in the thread context.
-     * The span must be manually closed to properly end it and revert the thread context.
-     *
-     * @param name the name of the span to be created
-     * @return an instance of {@link SpanScope}, which represents the created span and its associated context
-     */
+
+    public TracingOpenTelemetry(
+            Tracer tracer,
+            TextMapPropagator propagator) {
+
+        this.tracer = Objects.requireNonNull(
+                tracer,
+                "tracer must not be null"
+        );
+        this.propagator = Objects.requireNonNull(
+                propagator,
+                "propagator must not be null"
+        );
+    }
+
+
     public SpanScope addSpan(String name) {
-        Span span = tracer
-                .spanBuilder(name)
+        return addSpan(name, SpanKind.INTERNAL);
+    }
+
+
+    public SpanScope addSpan(
+            String name,
+            SpanKind kind) {
+
+        return addSpan(name, kind, Attributes.empty());
+    }
+
+
+    public SpanScope addSpan(
+            String name,
+            SpanKind kind,
+            Attributes attributes) {
+
+        Objects.requireNonNull(name, "name must not be null");
+        Objects.requireNonNull(kind, "kind must not be null");
+        Objects.requireNonNull(attributes, "attributes must not be null");
+
+        Span span = tracer.spanBuilder(name)
+                .setSpanKind(kind)
+                .setAllAttributes(attributes)
                 .startSpan();
 
         return new SpanScope(span);
     }
 
-    /**
-     * Retrieves the current active span in the execution context.
-     *
-     * @return the current {@link Span} if one is active, or a default no-op {@link Span} if none is active
-     */
+    public SpanScope addSpan(
+            String name,
+            SpanKind kind,
+            Attributes attributes,
+            Context parentContext) {
+
+        Objects.requireNonNull(parentContext, "parentContext must not be null");
+
+        Span span = tracer.spanBuilder(name)
+                .setParent(parentContext)
+                .setSpanKind(kind)
+                .setAllAttributes(attributes)
+                .startSpan();
+
+        return new SpanScope(span);
+    }
+
+
     public Span currentSpan() {
         return Span.current();
     }
 
-    /**
-     * Executes the specified operation within the context of a new span.
-     * The span is automatically managed and closed when the operation completes
-     * or an exception is thrown.
-     *
-     * @param name      the name of the span to be created
-     * @param operation the operation to be executed within the span's context
-     * @throws Exception if the provided operation throws an exception during execution
-     */
-    public <T> T withSpan(String name, SpanOperation<T> operation) throws Exception {
+
+    public <T> T withSpan(
+            String name,
+            SpanOperation<T> operation) throws Exception {
+
+        Objects.requireNonNull(operation, "operation must not be null");
+
         try (SpanScope scope = addSpan(name)) {
             try {
                 return operation.execute(scope.span());
-            } catch (Exception e) {
-                scope.recordException(e);
-                throw e;
+            } catch (Exception exception) {
+                scope.recordException(exception);
+                throw exception;
             }
         }
     }
 
-    /**
-     * Creates a new tracing instance using the global OpenTelemetry tracer.
-     *
-     * @return a new tracing instance
-     */
+    public <T> T captureLambdaHandler(
+            String name,
+            com.amazonaws.services.lambda.runtime.Context lambdaContext,
+            io.opentelemetry.context.Context parentContext,
+            SpanOperation<T> operation) throws Exception {
+
+        Objects.requireNonNull(name, "name must not be null");
+        Objects.requireNonNull(parentContext, "parentContext must not be null");
+        Objects.requireNonNull(operation, "operation must not be null");
+
+        Span span = tracer.spanBuilder(name)
+                .setParent(parentContext)
+                .setSpanKind(SpanKind.SERVER)
+                .setAttribute("faas.coldstart", LambdaHandlerProcessor.isColdStart())
+                .setAttribute("faas.invocation_id", lambdaContext.getAwsRequestId())
+                .startSpan();
+
+        try (SpanScope scope = new SpanScope(span)) {
+            try {
+                T result = operation.execute(span);
+
+                LambdaHandlerProcessor.coldStartDone();
+
+                return result;
+            } catch (Exception exception) {
+                scope.recordException(exception);
+                throw exception;
+            }
+        }
+    }
+
+    public <T> Context extractContext(
+            T carrier,
+            TextMapGetter<T> getter) {
+
+        return extractContext(Context.current(), carrier, getter);
+    }
+
+
+    public <T> Context extractContext(
+            Context context,
+            T carrier,
+            TextMapGetter<T> getter) {
+
+        Objects.requireNonNull(context, "context must not be null");
+        Objects.requireNonNull(getter, "getter must not be null");
+
+        return propagator.extract(
+                context,
+                carrier,
+                getter
+        );
+    }
+
+
+    public <T> void injectContext(
+            T carrier,
+            TextMapSetter<T> setter) {
+
+        injectContext(Context.current(), carrier, setter);
+    }
+
+
+    public <T> void injectContext(
+            Context context,
+            T carrier,
+            TextMapSetter<T> setter) {
+
+        Objects.requireNonNull(context, "context must not be null");
+        Objects.requireNonNull(setter, "setter must not be null");
+
+        propagator.inject(
+                context,
+                carrier,
+                setter
+        );
+    }
+
+    private static TextMapPropagator createDefaultPropagator() {
+        return OpenTelemetryProvider.propagator();
+    }
+
+
     public static TracingOpenTelemetry create() {
         return new TracingOpenTelemetry();
     }
 
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    public static final class Builder {
+
+        private Tracer tracer;
+        private TextMapPropagator propagator = createDefaultPropagator();
+
+        public Builder tracer(Tracer tracer) {
+            this.tracer = tracer;
+            return this;
+        }
+
+        public Builder propagator(TextMapPropagator propagator) {
+            this.propagator = propagator;
+            return this;
+        }
+
+        public TracingOpenTelemetry build() {
+            return new TracingOpenTelemetry(this);
+        }
+    }
 }