diff --git a/instrumentation/opentelemetry-instrumentation-api/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/SpanKeyInstrumentation.java b/instrumentation/opentelemetry-instrumentation-api/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/SpanKeyInstrumentation.java index 2439bfb7cc82..b8cc4ea99f40 100644 --- a/instrumentation/opentelemetry-instrumentation-api/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/SpanKeyInstrumentation.java +++ b/instrumentation/opentelemetry-instrumentation-api/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/SpanKeyInstrumentation.java @@ -41,22 +41,16 @@ public void transform(TypeTransformer transformer) { @SuppressWarnings("unused") public static class StoreInContextAdvice { - @Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class) - public static Object onEnter() { - return null; - } - - @Advice.OnMethodExit(suppress = Throwable.class) - public static void onExit( + @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class) + public static Context onEnter( @Advice.This SpanKey applicationSpanKey, @Advice.Argument(0) Context applicationContext, - @Advice.Argument(1) Span applicationSpan, - @Advice.Return(readOnly = false) Context newApplicationContext) { + @Advice.Argument(1) Span applicationSpan) { io.opentelemetry.instrumentation.api.internal.SpanKey agentSpanKey = SpanKeyBridging.toAgentOrNull(applicationSpanKey); if (agentSpanKey == null) { - return; + return null; } io.opentelemetry.context.Context agentContext = @@ -64,41 +58,61 @@ public static void onExit( io.opentelemetry.api.trace.Span agentSpan = Bridging.toAgentOrNull(applicationSpan); if (agentSpan == null) { - return; + // if application span can not be bridged to agent span, this could happen when it is not + // created through bridged GlobalOpenTelemetry, we'll let the original method run and + // store the span in context without bridging + return null; } io.opentelemetry.context.Context newAgentContext = agentSpanKey.storeInContext(agentContext, agentSpan); - newApplicationContext = AgentContextStorage.toApplicationContext(newAgentContext); + return AgentContextStorage.toApplicationContext(newAgentContext); + } + + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit( + @Advice.Enter Context newApplicationContext, + @Advice.Return(readOnly = false) Context result) { + + if (newApplicationContext != null) { + result = newApplicationContext; + } } } @SuppressWarnings("unused") public static class FromContextOrNullAdvice { - @Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class) - public static Object onEnter() { - return null; - } - - @Advice.OnMethodExit(suppress = Throwable.class) - public static void onExit( - @Advice.This SpanKey applicationSpanKey, - @Advice.Argument(0) Context applicationContext, - @Advice.Return(readOnly = false) Span applicationSpan) { + @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class) + public static Span onEnter( + @Advice.This SpanKey applicationSpanKey, @Advice.Argument(0) Context applicationContext) { io.opentelemetry.instrumentation.api.internal.SpanKey agentSpanKey = SpanKeyBridging.toAgentOrNull(applicationSpanKey); if (agentSpanKey == null) { - return; + return null; } io.opentelemetry.context.Context agentContext = AgentContextStorage.getAgentContext(applicationContext); io.opentelemetry.api.trace.Span agentSpan = agentSpanKey.fromContextOrNull(agentContext); + if (agentSpan == null) { + // Bridged agent span was not found. Run the original method, there could be an unbridged + // span stored in the application context. + return null; + } - applicationSpan = agentSpan == null ? null : Bridging.toApplication(agentSpan); + return Bridging.toApplication(agentSpan); + } + + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit( + @Advice.Enter Span applicationSpan, @Advice.Return(readOnly = false) Span result) { + + if (applicationSpan != null) { + result = applicationSpan; + } } } } diff --git a/instrumentation/opentelemetry-instrumentation-api/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/ContextBridgeTest.java b/instrumentation/opentelemetry-instrumentation-api/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/ContextBridgeTest.java index cd5078cd2211..2ebea20102b6 100644 --- a/instrumentation/opentelemetry-instrumentation-api/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/ContextBridgeTest.java +++ b/instrumentation/opentelemetry-instrumentation-api/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/ContextBridgeTest.java @@ -6,6 +6,7 @@ package io.opentelemetry.javaagent.instrumentation.instrumentationapi; import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo; +import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.Assertions.assertNotNull; import io.opentelemetry.api.trace.Span; @@ -18,6 +19,7 @@ import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension; import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension; import io.opentelemetry.javaagent.instrumentation.testing.AgentSpanTesting; +import io.opentelemetry.sdk.OpenTelemetrySdk; import io.opentelemetry.semconv.ErrorAttributes; import io.opentelemetry.semconv.HttpAttributes; import java.util.Arrays; @@ -76,6 +78,18 @@ void testSpanKeyBridge() { }); } + @Test + void testSpanKeyBridge_UnbridgedSpan() { + OpenTelemetrySdk openTelemetry = OpenTelemetrySdk.builder().build(); + // span is bridged only when it is created though a bridged OpenTelemetry instance obtained + // from GlobalOpenTelemetry + Span span = openTelemetry.getTracer("test").spanBuilder("test").startSpan(); + + Context context = SpanKey.HTTP_CLIENT.storeInContext(Context.current(), span); + assertThat(context).isNotNull(); + assertThat(SpanKey.HTTP_CLIENT.fromContextOrNull(context)).isEqualTo(span); + } + @Test void testHttpRouteHolder_SameSourceAsServerInstrumentationDoesNotOverrideRoute() { AgentSpanTesting.runWithHttpServerSpan(