Skip to content

Commit

Permalink
Bridge span keys defined in instrumentation api (#3911)
Browse files Browse the repository at this point in the history
  • Loading branch information
laurit committed Aug 26, 2021
1 parent 28db0c3 commit b5aec6a
Show file tree
Hide file tree
Showing 10 changed files with 390 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ public final class SpanKey {
private static final ContextKey<Span> CONSUMER_KEY =
ContextKey.named("opentelemetry-traces-span-key-consumer");

// TODO bridge these constants in AgentContextStorage
private static final ContextKey<Span> HTTP_KEY =
ContextKey.named("opentelemetry-traces-span-key-http");
private static final ContextKey<Span> RPC_KEY =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,12 @@ dependencies {
// using OpenTelemetry SDK to make sure that instrumentation doesn't cause
// OpenTelemetrySdk.getTracerProvider() to throw ClassCastException
testImplementation("io.opentelemetry:opentelemetry-sdk")
testImplementation(project(":instrumentation-api"))

// @WithSpan annotation is used to generate spans in ContextBridgeTest
testImplementation("io.opentelemetry:opentelemetry-extension-annotations")
testInstrumentation(project(":instrumentation:opentelemetry-annotations-1.0:javaagent"))

testImplementation(project(":instrumentation:opentelemetry-api-1.0:testing"))
testInstrumentation(project(":instrumentation:opentelemetry-api-1.0:testing"))
}
Original file line number Diff line number Diff line change
Expand Up @@ -151,8 +151,26 @@ public static Context newContextWrapper(
"io.opentelemetry.api.baggage.BaggageContextKey",
BaggageBridging::toApplication,
BaggageBridging::toAgent),
bridgeSpanKey("SERVER_KEY"),
bridgeSpanKey("CONSUMER_KEY"),
bridgeSpanKey("HTTP_KEY"),
bridgeSpanKey("RPC_KEY"),
bridgeSpanKey("DB_KEY"),
bridgeSpanKey("MESSAGING_KEY"),
bridgeSpanKey("CLIENT_KEY"),
bridgeSpanKey("PRODUCER_KEY"),
};

private static ContextKeyBridge<Span, io.opentelemetry.api.trace.Span> bridgeSpanKey(
String name) {
return new ContextKeyBridge<>(
"application.io.opentelemetry.instrumentation.api.instrumenter.SpanKey",
"io.opentelemetry.instrumentation.api.instrumenter.SpanKey",
name,
Bridging::toApplication,
Bridging::toAgentOrNull);
}

@Override
public Scope attach(Context toAttach) {
io.opentelemetry.context.Context currentAgentContext =
Expand Down Expand Up @@ -266,13 +284,22 @@ static class ContextKeyBridge<APPLICATION, AGENT> {
String agentKeyHolderClassName,
Function<AGENT, APPLICATION> toApplication,
Function<APPLICATION, AGENT> toAgent) {
this(applicationKeyHolderClassName, agentKeyHolderClassName, "KEY", toApplication, toAgent);
}

ContextKeyBridge(
String applicationKeyHolderClassName,
String agentKeyHolderClassName,
String fieldName,
Function<AGENT, APPLICATION> toApplication,
Function<APPLICATION, AGENT> toAgent) {
this.toApplication = toApplication;
this.toAgent = toAgent;

ContextKey<APPLICATION> applicationContextKey;
try {
Class<?> applicationKeyHolderClass = Class.forName(applicationKeyHolderClassName);
Field applicationContextKeyField = applicationKeyHolderClass.getDeclaredField("KEY");
Field applicationContextKeyField = applicationKeyHolderClass.getDeclaredField(fieldName);
applicationContextKeyField.setAccessible(true);
applicationContextKey = (ContextKey<APPLICATION>) applicationContextKeyField.get(null);
} catch (Throwable t) {
Expand All @@ -283,7 +310,7 @@ static class ContextKeyBridge<APPLICATION, AGENT> {
io.opentelemetry.context.ContextKey<AGENT> agentContextKey;
try {
Class<?> agentKeyHolderClass = Class.forName(agentKeyHolderClassName);
Field agentContextKeyField = agentKeyHolderClass.getDeclaredField("KEY");
Field agentContextKeyField = agentKeyHolderClass.getDeclaredField(fieldName);
agentContextKeyField.setAccessible(true);
agentContextKey =
(io.opentelemetry.context.ContextKey<AGENT>) agentContextKeyField.get(null);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@ import io.opentelemetry.api.trace.Span
import io.opentelemetry.context.Context
import io.opentelemetry.context.ContextKey
import io.opentelemetry.extension.annotations.WithSpan
import io.opentelemetry.instrumentation.api.instrumenter.SpanKey
import io.opentelemetry.instrumentation.api.tracer.ClientSpan
import io.opentelemetry.instrumentation.api.tracer.ConsumerSpan
import io.opentelemetry.instrumentation.api.tracer.ServerSpan
import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
import java.util.concurrent.CountDownLatch
import java.util.concurrent.Executors
Expand Down Expand Up @@ -145,6 +149,64 @@ class ContextBridgeTest extends AgentInstrumentationSpecification {
Context.current() == Context.root()
}

def "test server span bridge"() {
expect:
AgentSpanTesting.runWithServerSpan("server") {
assert Span.current() != null
assert ServerSpan.fromContextOrNull(Context.current()) != null
runWithSpan("internal") {
assert ServerSpan.fromContextOrNull(Context.current()) != null
}
}
}

def "test consumer span bridge"() {
expect:
AgentSpanTesting.runWithConsumerSpan("consumer") {
assert Span.current() != null
assert ConsumerSpan.fromContextOrNull(Context.current()) != null
runWithSpan("internal") {
assert ConsumerSpan.fromContextOrNull(Context.current()) != null
}
}
}

def "test client span bridge"() {
expect:
AgentSpanTesting.runWithClientSpan("client") {
assert Span.current() != null
assert ClientSpan.fromContextOrNull(Context.current()) != null
runWithSpan("internal") {
assert ClientSpan.fromContextOrNull(Context.current()) != null
}
}
}

def "test span key bridge"() {
expect:
AgentSpanTesting.runWithAllSpanKeys("parent") {
assert Span.current() != null
def spanKeys = [
SpanKey.SERVER,
SpanKey.CONSUMER,
SpanKey.HTTP_CLIENT,
SpanKey.RPC_CLIENT,
SpanKey.DB_CLIENT,
SpanKey.MESSAGING_PRODUCER,
SpanKey.ALL_CLIENTS,
SpanKey.ALL_PRODUCERS
]
spanKeys.each { spanKey ->
assert spanKey.fromContextOrNull(Context.current()) != null
}
runWithSpan("internal") {
spanKeys.each { spanKey ->
assert spanKey.fromContextOrNull(Context.current()) != null
}
}
}
}

// TODO (trask)
// more tests are needed here, not sure how to implement, probably need to write some test
// instrumentation to help test, similar to :testing-common:integration-tests
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
plugins {
id("otel.javaagent-testing")
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/

public class AgentSpanTesting {

/**
* Runs the provided {@code runnable} inside the scope of an SERVER span with name {@code
* spanName}.
*/
public static void runWithServerSpan(String spanName, Runnable runnable) {
runnable.run();
}

/**
* Runs the provided {@code runnable} inside the scope of an CONSUMER span with name {@code
* spanName}.
*/
public static void runWithConsumerSpan(String spanName, Runnable runnable) {
runnable.run();
}

/**
* Runs the provided {@code runnable} inside the scope of an CLIENT span with name {@code
* spanName}.
*/
public static void runWithClientSpan(String spanName, Runnable runnable) {
runnable.run();
}

/**
* Runs the provided {@code runnable} inside the scope of an INTERNAL span with name {@code
* spanName}. Span is added into context under all possible keys from {@link
* io.opentelemetry.instrumentation.api.instrumenter.SpanKey}
*/
public static void runWithAllSpanKeys(String spanName, Runnable runnable) {
runnable.run();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,138 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/

import static net.bytebuddy.matcher.ElementMatchers.named;

import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

public class AgentSpanTestingInstrumentation implements TypeInstrumentation {

@Override
public ElementMatcher<TypeDescription> typeMatcher() {
return named("AgentSpanTesting");
}

@Override
public void transform(TypeTransformer transformer) {
transformer.applyAdviceToMethod(
named("runWithServerSpan"), this.getClass().getName() + "$RunWithServerSpanAdvice");
transformer.applyAdviceToMethod(
named("runWithConsumerSpan"), this.getClass().getName() + "$RunWithConsumerSpanAdvice");
transformer.applyAdviceToMethod(
named("runWithClientSpan"), this.getClass().getName() + "$RunWithClientSpanAdvice");
transformer.applyAdviceToMethod(
named("runWithAllSpanKeys"), this.getClass().getName() + "$RunWithAllSpanKeysAdvice");
}

@SuppressWarnings("unused")
public static class RunWithServerSpanAdvice {

@Advice.OnMethodEnter(suppress = Throwable.class)
public static void onEnter(
@Advice.Argument(0) String spanName,
@Advice.Local("otelContext") Context context,
@Advice.Local("otelScope") Scope scope) {
context = AgentSpanTestingTracer.tracer().startServerSpan(spanName);
scope = context.makeCurrent();
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void onExit(
@Advice.Thrown Throwable throwable,
@Advice.Local("otelContext") Context context,
@Advice.Local("otelScope") Scope scope) {
scope.close();
if (throwable != null) {
AgentSpanTestingTracer.tracer().endExceptionally(context, throwable);
} else {
AgentSpanTestingTracer.tracer().end(context);
}
}
}

@SuppressWarnings("unused")
public static class RunWithConsumerSpanAdvice {

@Advice.OnMethodEnter(suppress = Throwable.class)
public static void onEnter(
@Advice.Argument(0) String spanName,
@Advice.Local("otelContext") Context context,
@Advice.Local("otelScope") Scope scope) {
context = AgentSpanTestingTracer.tracer().startConsumerSpan(spanName);
scope = context.makeCurrent();
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void onExit(
@Advice.Thrown Throwable throwable,
@Advice.Local("otelContext") Context context,
@Advice.Local("otelScope") Scope scope) {
scope.close();
if (throwable != null) {
AgentSpanTestingTracer.tracer().endExceptionally(context, throwable);
} else {
AgentSpanTestingTracer.tracer().end(context);
}
}
}

@SuppressWarnings("unused")
public static class RunWithClientSpanAdvice {

@Advice.OnMethodEnter(suppress = Throwable.class)
public static void onEnter(
@Advice.Argument(0) String spanName,
@Advice.Local("otelContext") Context context,
@Advice.Local("otelScope") Scope scope) {
context = AgentSpanTestingTracer.tracer().startClientSpan(spanName);
scope = context.makeCurrent();
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void onExit(
@Advice.Thrown Throwable throwable,
@Advice.Local("otelContext") Context context,
@Advice.Local("otelScope") Scope scope) {
scope.close();
if (throwable != null) {
AgentSpanTestingTracer.tracer().endExceptionally(context, throwable);
} else {
AgentSpanTestingTracer.tracer().end(context);
}
}
}

@SuppressWarnings("unused")
public static class RunWithAllSpanKeysAdvice {

@Advice.OnMethodEnter(suppress = Throwable.class)
public static void onEnter(
@Advice.Argument(0) String spanName,
@Advice.Local("otelContext") Context context,
@Advice.Local("otelScope") Scope scope) {
context = AgentSpanTestingTracer.tracer().startSpanWithAllKeys(spanName);
scope = context.makeCurrent();
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void onExit(
@Advice.Thrown Throwable throwable,
@Advice.Local("otelContext") Context context,
@Advice.Local("otelScope") Scope scope) {
scope.close();
if (throwable != null) {
AgentSpanTestingTracer.tracer().endExceptionally(context, throwable);
} else {
AgentSpanTestingTracer.tracer().end(context);
}
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/

import static java.util.Collections.singletonList;

import com.google.auto.service.AutoService;
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import java.util.List;

@AutoService(InstrumentationModule.class)
public class AgentSpanTestingInstrumentationModule extends InstrumentationModule {
public AgentSpanTestingInstrumentationModule() {
super(AgentSpanTestingInstrumentationModule.class.getName());
}

@Override
public boolean isHelperClass(String className) {
return className.startsWith("AgentSpanTestingTracer");
}

@Override
public List<TypeInstrumentation> typeInstrumentations() {
return singletonList(new AgentSpanTestingInstrumentation());
}
}
Loading

0 comments on commit b5aec6a

Please sign in to comment.