Refactor AWS lambda tracers so that they use Context
... and do not use BaseTracer#startScope(). This PR eliminates all overrides of startScope(), in the next one I'll remove it.
MODIFY
instrumentation/aws-lambda-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awslambda/v1_0/AwsLambdaRequestHandlerInstrumentation.java
instrumentation/aws-lambda-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awslambda/v1_0/AwsLambdaRequestHandlerInstrumentation.java
@@ -16,7 +16,6 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.events.SQSEvent;
-import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Span.Kind;
import io.opentelemetry.context.Scope;
import io.opentelemetry.javaagent.instrumentation.api.OpenTelemetrySdkAccess;
@@ -57,40 +56,40 @@ public class AwsLambdaRequestHandlerInstrumentation implements TypeInstrumentati
public static void onEnter(
@Advice.Argument(value = 0, typing = Typing.DYNAMIC) Object arg,
@Advice.Argument(1) Context context,
- @Advice.Local("otelFunctionSpan") Span functionSpan,
+ @Advice.Local("otelFunctionContext") io.opentelemetry.context.Context functionContext,
@Advice.Local("otelFunctionScope") Scope functionScope,
- @Advice.Local("otelMessageSpan") Span messageSpan,
+ @Advice.Local("otelMessageContext") io.opentelemetry.context.Context messageContext,
@Advice.Local("otelMessageScope") Scope messageScope) {
- functionSpan = functionTracer().startSpan(context, arg, Kind.SERVER);
- functionScope = functionTracer().startScope(functionSpan);
+ functionContext = functionTracer().startSpan(context, Kind.SERVER, arg);
+ functionScope = functionContext.makeCurrent();
if (arg instanceof SQSEvent) {
- messageSpan = messageTracer().startSpan(context, (SQSEvent) arg);
- messageScope = messageTracer().startScope(messageSpan);
+ messageContext = messageTracer().startSpan((SQSEvent) arg);
+ messageScope = messageContext.makeCurrent();
}
}
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void stopSpan(
@Advice.Thrown Throwable throwable,
- @Advice.Local("otelFunctionSpan") Span functionSpan,
+ @Advice.Local("otelFunctionContext") io.opentelemetry.context.Context functionContext,
@Advice.Local("otelFunctionScope") Scope functionScope,
- @Advice.Local("otelMessageSpan") Span messageSpan,
+ @Advice.Local("otelMessageContext") io.opentelemetry.context.Context messageContext,
@Advice.Local("otelMessageScope") Scope messageScope) {
if (messageScope != null) {
messageScope.close();
if (throwable != null) {
- messageTracer().endExceptionally(messageSpan, throwable);
+ messageTracer().endExceptionally(messageContext, throwable);
} else {
- messageTracer().end(messageSpan);
+ messageTracer().end(messageContext);
}
}
functionScope.close();
if (throwable != null) {
- functionTracer().endExceptionally(functionSpan, throwable);
+ functionTracer().endExceptionally(functionContext, throwable);
} else {
- functionTracer().end(functionSpan);
+ functionTracer().end(functionContext);
}
OpenTelemetrySdkAccess.forceFlush(1, TimeUnit.SECONDS);
}

MODIFY
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/AwsLambdaMessageTracer.java
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/AwsLambdaMessageTracer.java
@@ -12,7 +12,7 @@ import io.opentelemetry.api.trace.Span.Kind;
import io.opentelemetry.api.trace.SpanBuilder;
import io.opentelemetry.api.trace.SpanContext;
import io.opentelemetry.api.trace.Tracer;
-import io.opentelemetry.context.Scope;
+import io.opentelemetry.context.Context;
import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
@@ -26,7 +26,7 @@ public class AwsLambdaMessageTracer extends BaseTracer {
super(tracer);
}
- public Span startSpan(com.amazonaws.services.lambda.runtime.Context context, SQSEvent event) {
+ public Context startSpan(SQSEvent event) {
// Use event source in name if all messages have the same source, otherwise use placeholder.
String source = "multiple_sources";
if (!event.getRecords().isEmpty()) {
@@ -52,10 +52,10 @@ public class AwsLambdaMessageTracer extends BaseTracer {
addLinkToMessageParent(message, span);
}
- return span.startSpan();
+ return Context.current().with(span.startSpan());
}
- public Span startSpan(SQSMessage message) {
+ public Context startSpan(SQSMessage message) {
SpanBuilder span =
tracer.spanBuilder(message.getEventSource() + " process").setSpanKind(Kind.CONSUMER);
@@ -66,12 +66,7 @@ public class AwsLambdaMessageTracer extends BaseTracer {
addLinkToMessageParent(message, span);
- return span.startSpan();
- }
-
- @Override
- public Scope startScope(Span span) {
- return span.makeCurrent();
+ return Context.current().with(span.startSpan());
}
private void addLinkToMessageParent(SQSMessage message, SpanBuilder span) {

MODIFY
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/AwsLambdaTracer.java
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/AwsLambdaTracer.java
@@ -16,9 +16,7 @@ import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Span.Kind;
import io.opentelemetry.api.trace.SpanBuilder;
-import io.opentelemetry.api.trace.SpanContext;
import io.opentelemetry.api.trace.Tracer;
-import io.opentelemetry.context.Scope;
import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
import io.opentelemetry.semconv.trace.attributes.SemanticAttributes.FaasTriggerValues;
import java.lang.invoke.MethodHandle;
@@ -30,7 +28,6 @@ import org.checkerframework.checker.nullness.qual.Nullable;
public class AwsLambdaTracer extends BaseTracer {
- private static final String AWS_TRACE_HEADER_ENV_KEY = "_X_AMZN_TRACE_ID";
private static final MethodHandle GET_FUNCTION_ARN;
static {
@@ -56,40 +53,6 @@ public class AwsLambdaTracer extends BaseTracer {
super(tracer);
}
- private boolean isValid(io.opentelemetry.context.Context context) {
- if (context == null) {
- return false;
- }
- Span parentSpan = Span.fromContext(context);
- SpanContext parentSpanContext = parentSpan.getSpanContext();
- return parentSpanContext.isValid();
- }
-
- private io.opentelemetry.context.Context parent(Map<String, String> headers) {
-
- io.opentelemetry.context.Context parentContext = null;
- String parentTraceHeader = System.getenv(AWS_TRACE_HEADER_ENV_KEY);
- if (parentTraceHeader != null) {
- parentContext = ParentContextExtractor.fromXRayHeader(parentTraceHeader);
- }
- if (!isValid(parentContext)) {
- // try http
- parentContext = ParentContextExtractor.fromHttpHeaders(headers);
- }
-
- return parentContext;
- }
-
- private SpanBuilder createSpan(Context context, Object input, Map<String, String> headers) {
- SpanBuilder span = tracer.spanBuilder(spanName(context, input));
- setAttributes(span, context, input);
- io.opentelemetry.context.Context parent = parent(headers);
- if (parent != null) {
- span.setParent(parent);
- }
- return span;
- }
-
private void setAttributes(SpanBuilder span, Context context, Object input) {
setCommonAttributes(span, context);
if (input instanceof APIGatewayProxyRequestEvent) {
@@ -148,25 +111,25 @@ public class AwsLambdaTracer extends BaseTracer {
return name == null ? context.getFunctionName() : name;
}
- public Span startSpan(Context context, Kind kind, Object input, Map<String, String> headers) {
- return createSpan(context, input, headers).setSpanKind(kind).startSpan();
+ public io.opentelemetry.context.Context startSpan(Context awsContext, Kind kind, Object input) {
+ return startSpan(awsContext, kind, input, Collections.emptyMap());
}
- public Span startSpan(Context context, Object input, Kind kind) {
- return createSpan(context, input, Collections.emptyMap()).setSpanKind(kind).startSpan();
- }
+ public io.opentelemetry.context.Context startSpan(
+ Context awsContext, Kind kind, Object input, Map<String, String> headers) {
+ io.opentelemetry.context.Context parentContext = ParentContextExtractor.extract(headers);
- /** Creates new scoped context with the given span. */
- @Override
- public Scope startScope(Span span) {
- io.opentelemetry.context.Context newContext =
- withServerSpan(io.opentelemetry.context.Context.current(), span);
- return newContext.makeCurrent();
+ SpanBuilder spanBuilder = tracer.spanBuilder(spanName(awsContext, input));
+ setAttributes(spanBuilder, awsContext, input);
+ Span span = spanBuilder.setParent(parentContext).setSpanKind(kind).startSpan();
+
+ return withServerSpan(parentContext, span);
}
- public void onOutput(Span span, Object output) {
+ public void onOutput(io.opentelemetry.context.Context context, Object output) {
if (output instanceof APIGatewayProxyResponseEvent) {
- httpSpanAttributes.onResponse(span, (APIGatewayProxyResponseEvent) output);
+ httpSpanAttributes.onResponse(
+ Span.fromContext(context), (APIGatewayProxyResponseEvent) output);
}
}

MODIFY
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/ParentContextExtractor.java
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/ParentContextExtractor.java
@@ -7,20 +7,43 @@ package io.opentelemetry.instrumentation.awslambda.v1_0;
import static io.opentelemetry.instrumentation.awslambda.v1_0.MapUtils.lowercaseMap;
-import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.propagation.TextMapPropagator;
import io.opentelemetry.extension.trace.propagation.AwsXrayPropagator;
+import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
import java.util.Collections;
import java.util.Map;
public class ParentContextExtractor {
+ private static final String AWS_TRACE_HEADER_ENV_KEY = "_X_AMZN_TRACE_ID";
+
+ static Context extract(Map<String, String> headers) {
+ Context parentContext = null;
+ String parentTraceHeader = System.getenv(AWS_TRACE_HEADER_ENV_KEY);
+ if (parentTraceHeader != null) {
+ parentContext = ParentContextExtractor.fromXRayHeader(parentTraceHeader);
+ }
+ if (!isValid(parentContext)) {
+ // try http
+ parentContext = ParentContextExtractor.fromHttpHeaders(headers);
+ }
+ return parentContext;
+ }
+
+ private static boolean isValid(Context context) {
+ if (context == null) {
+ return false;
+ }
+ Span parentSpan = Span.fromContext(context);
+ SpanContext parentSpanContext = parentSpan.getSpanContext();
+ return parentSpanContext.isValid();
+ }
+
static Context fromHttpHeaders(Map<String, String> headers) {
- return GlobalOpenTelemetry.getPropagators()
- .getTextMapPropagator()
- .extract(
- io.opentelemetry.context.Context.current(), lowercaseMap(headers), MapGetter.INSTANCE);
+ return BaseTracer.extract(lowercaseMap(headers), MapGetter.INSTANCE);
}
// lower-case map getter used for extraction
@@ -29,7 +52,8 @@ public class ParentContextExtractor {
static Context fromXRayHeader(String parentHeader) {
return AwsXrayPropagator.getInstance()
.extract(
- Context.current(),
+ // see BaseTracer#extract() on why we're using root() here
+ Context.root(),
Collections.singletonMap(AWS_TRACE_HEADER_PROPAGATOR_KEY, parentHeader),
MapGetter.INSTANCE);
}

MODIFY
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingRequestHandler.java
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingRequestHandler.java
@@ -8,7 +8,6 @@ package io.opentelemetry.instrumentation.awslambda.v1_0;
import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.RequestHandler;
import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
-import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Span.Kind;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Scope;
@@ -68,20 +67,21 @@ public abstract class TracingRequestHandler<I, O> implements RequestHandler<I, O
@Override
public final O handleRequest(I input, Context context) {
- Span span = tracer.startSpan(context, Kind.SERVER, input, getHeaders(input));
+ io.opentelemetry.context.Context otelContext =
+ tracer.startSpan(context, Kind.SERVER, input, getHeaders(input));
Throwable error = null;
- try (Scope ignored = tracer.startScope(span)) {
+ try (Scope ignored = otelContext.makeCurrent()) {
O output = doHandleRequest(input, context);
- tracer.onOutput(span, output);
+ tracer.onOutput(otelContext, output);
return output;
} catch (Throwable t) {
error = t;
throw t;
} finally {
if (error != null) {
- tracer.endExceptionally(span, error);
+ tracer.endExceptionally(otelContext, error);
} else {
- tracer.end(span);
+ tracer.end(otelContext);
}
OpenTelemetrySdk.getGlobalTracerManagement()
.forceFlush()

MODIFY
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingRequestStreamHandler.java
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingRequestStreamHandler.java
@@ -7,7 +7,6 @@ package io.opentelemetry.instrumentation.awslambda.v1_0;
import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
-import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Span.Kind;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Scope;
@@ -69,12 +68,14 @@ public abstract class TracingRequestStreamHandler implements RequestStreamHandle
throws IOException {
ApiGatewayProxyRequest proxyRequest = ApiGatewayProxyRequest.forStream(input);
- Span span = tracer.startSpan(context, Kind.SERVER, input, proxyRequest.getHeaders());
+ io.opentelemetry.context.Context otelContext =
+ tracer.startSpan(context, Kind.SERVER, input, proxyRequest.getHeaders());
- try (Scope ignored = tracer.startScope(span)) {
- doHandleRequest(proxyRequest.freshStream(), new OutputStreamWrapper(output, span), context);
+ try (Scope ignored = otelContext.makeCurrent()) {
+ doHandleRequest(
+ proxyRequest.freshStream(), new OutputStreamWrapper(output, otelContext), context);
} catch (Throwable t) {
- tracer.endExceptionally(span, t);
+ tracer.endExceptionally(otelContext, t);
OpenTelemetrySdk.getGlobalTracerManagement()
.forceFlush()
.join(flushTimeout, TimeUnit.SECONDS);
@@ -88,11 +89,12 @@ public abstract class TracingRequestStreamHandler implements RequestStreamHandle
private class OutputStreamWrapper extends OutputStream {
private final OutputStream delegate;
- private final Span span;
+ private final io.opentelemetry.context.Context otelContext;
- OutputStreamWrapper(OutputStream delegate, Span span) {
+ private OutputStreamWrapper(
+ OutputStream delegate, io.opentelemetry.context.Context otelContext) {
this.delegate = delegate;
- this.span = span;
+ this.otelContext = otelContext;
}
@Override
@@ -118,7 +120,7 @@ public abstract class TracingRequestStreamHandler implements RequestStreamHandle
@Override
public void close() throws IOException {
delegate.close();
- tracer.end(span);
+ tracer.end(otelContext);
OpenTelemetrySdk.getGlobalTracerManagement().forceFlush().join(1, TimeUnit.SECONDS);
}
}

MODIFY
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingSqsEventHandler.java
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingSqsEventHandler.java
@@ -7,7 +7,6 @@ package io.opentelemetry.instrumentation.awslambda.v1_0;
import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.events.SQSEvent;
-import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Scope;
import io.opentelemetry.sdk.OpenTelemetrySdk;
@@ -40,18 +39,18 @@ public abstract class TracingSqsEventHandler extends TracingRequestHandler<SQSEv
@Override
public Void doHandleRequest(SQSEvent event, Context context) {
- Span span = tracer.startSpan(context, event);
+ io.opentelemetry.context.Context otelContext = tracer.startSpan(event);
Throwable error = null;
- try (Scope ignored = tracer.startScope(span)) {
+ try (Scope ignored = otelContext.makeCurrent()) {
handleEvent(event, context);
} catch (Throwable t) {
error = t;
throw t;
} finally {
if (error != null) {
- tracer.endExceptionally(span, error);
+ tracer.endExceptionally(otelContext, error);
} else {
- tracer.end(span);
+ tracer.end(otelContext);
}
OpenTelemetrySdk.getGlobalTracerManagement().forceFlush().join(1, TimeUnit.SECONDS);
}

MODIFY
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingSqsMessageHandler.java
instrumentation/aws-lambda-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambda/v1_0/TracingSqsMessageHandler.java
@@ -8,7 +8,6 @@ package io.opentelemetry.instrumentation.awslambda.v1_0;
import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.events.SQSEvent;
import com.amazonaws.services.lambda.runtime.events.SQSEvent.SQSMessage;
-import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Scope;
@@ -37,18 +36,18 @@ public abstract class TracingSqsMessageHandler extends TracingSqsEventHandler {
@Override
protected final void handleEvent(SQSEvent event, Context context) {
for (SQSMessage message : event.getRecords()) {
- Span span = getTracer().startSpan(message);
+ io.opentelemetry.context.Context otelContext = getTracer().startSpan(message);
Throwable error = null;
- try (Scope ignored = getTracer().startScope(span)) {
+ try (Scope ignored = otelContext.makeCurrent()) {
handleMessage(message, context);
} catch (Throwable t) {
error = t;
throw t;
} finally {
if (error != null) {
- getTracer().endExceptionally(span, error);
+ getTracer().endExceptionally(otelContext, error);
} else {
- getTracer().end(span);
+ getTracer().end(otelContext);
}
}
}

Refactor AWS lambda tracers so that they use Context (#2194)

... and do not use BaseTracer#startScope(). This PR eliminates all overrides of startScope(), in the next one I'll remove it.