diff --git a/gradle/enforcement/checkstyle.xml b/gradle/enforcement/checkstyle.xml index 656397805c..64dbb91b9e 100644 --- a/gradle/enforcement/checkstyle.xml +++ b/gradle/enforcement/checkstyle.xml @@ -227,7 +227,7 @@ value="Class type name ''{0}'' must match pattern ''{1}''."/> - + diff --git a/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/context/AgentContextStorage.java b/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/context/AgentContextStorage.java index 8df3a25bda..d0bac8a57e 100644 --- a/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/context/AgentContextStorage.java +++ b/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/context/AgentContextStorage.java @@ -14,6 +14,7 @@ import application.io.opentelemetry.context.Scope; import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.baggage.BaggageBridging; import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.trace.Bridging; import java.lang.reflect.Field; +import java.util.function.Function; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -59,66 +60,19 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable { static final io.opentelemetry.context.ContextKey APPLICATION_CONTEXT = io.opentelemetry.context.ContextKey.named("otel-context"); - static final io.opentelemetry.context.ContextKey - AGENT_SPAN_CONTEXT_KEY; - @Nullable static final ContextKey APPLICATION_SPAN_CONTEXT_KEY; - - static final io.opentelemetry.context.ContextKey - AGENT_BAGGAGE_CONTEXT_KEY; - @Nullable static final ContextKey APPLICATION_BAGGAGE_CONTEXT_KEY; - - static { - io.opentelemetry.context.ContextKey agentSpanContextKey; - try { - Class spanContextKey = Class.forName("io.opentelemetry.api.trace.SpanContextKey"); - Field spanContextKeyField = spanContextKey.getDeclaredField("KEY"); - spanContextKeyField.setAccessible(true); - agentSpanContextKey = - (io.opentelemetry.context.ContextKey) - spanContextKeyField.get(null); - } catch (Throwable t) { - agentSpanContextKey = null; - } - AGENT_SPAN_CONTEXT_KEY = agentSpanContextKey; - - ContextKey applicationSpanContextKey; - try { - Class spanContextKey = - Class.forName("application.io.opentelemetry.api.trace.SpanContextKey"); - Field spanContextKeyField = spanContextKey.getDeclaredField("KEY"); - spanContextKeyField.setAccessible(true); - applicationSpanContextKey = (ContextKey) spanContextKeyField.get(null); - } catch (Throwable t) { - applicationSpanContextKey = null; - } - APPLICATION_SPAN_CONTEXT_KEY = applicationSpanContextKey; - - io.opentelemetry.context.ContextKey - agentBaggageContextKey; - try { - Class baggageContextKey = Class.forName("io.opentelemetry.api.baggage.BaggageContextKey"); - Field baggageContextKeyField = baggageContextKey.getDeclaredField("KEY"); - baggageContextKeyField.setAccessible(true); - agentBaggageContextKey = - (io.opentelemetry.context.ContextKey) - baggageContextKeyField.get(null); - } catch (Throwable t) { - agentBaggageContextKey = null; - } - AGENT_BAGGAGE_CONTEXT_KEY = agentBaggageContextKey; - - ContextKey applicationBaggageContextKey; - try { - Class baggageContextKey = - Class.forName("application.io.opentelemetry.api.baggage.BaggageContextKey"); - Field baggageContextKeyField = baggageContextKey.getDeclaredField("KEY"); - baggageContextKeyField.setAccessible(true); - applicationBaggageContextKey = (ContextKey) baggageContextKeyField.get(null); - } catch (Throwable t) { - applicationBaggageContextKey = null; - } - APPLICATION_BAGGAGE_CONTEXT_KEY = applicationBaggageContextKey; - } + static final ContextKeyBridge[] CONTEXT_KEY_BRIDGES = + new ContextKeyBridge[] { + new ContextKeyBridge( + "application.io.opentelemetry.api.trace.SpanContextKey", + "io.opentelemetry.api.trace.SpanContextKey", + Bridging::toApplication, + Bridging::toAgentOrNull), + new ContextKeyBridge<>( + "application.io.opentelemetry.api.baggage.BaggageContextKey", + "io.opentelemetry.api.baggage.BaggageContextKey", + BaggageBridging::toApplication, + BaggageBridging::toAgent), + }; @Override public Scope attach(Context toAttach) { @@ -163,8 +117,8 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable { } private static class AgentContextWrapper implements Context { - private final io.opentelemetry.context.Context agentContext; - private final Context applicationContext; + final io.opentelemetry.context.Context agentContext; + final Context applicationContext; AgentContextWrapper(io.opentelemetry.context.Context agentContext, Context applicationContext) { this.agentContext = agentContext; @@ -180,47 +134,23 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable { @Override public V get(ContextKey key) { - if (key == APPLICATION_SPAN_CONTEXT_KEY) { - io.opentelemetry.api.trace.Span agentSpan = agentContext.get(AGENT_SPAN_CONTEXT_KEY); - if (agentSpan == null) { - return null; + for (ContextKeyBridge bridge : CONTEXT_KEY_BRIDGES) { + V value = bridge.get(this, key); + if (value != null) { + return value; } - Span applicationSpan = Bridging.toApplication(agentSpan); - @SuppressWarnings("unchecked") - V value = (V) applicationSpan; - return value; - } - if (key == APPLICATION_BAGGAGE_CONTEXT_KEY) { - io.opentelemetry.api.baggage.Baggage agentBaggage = - agentContext.get(AGENT_BAGGAGE_CONTEXT_KEY); - if (agentBaggage == null) { - return null; - } - Baggage applicationBaggage = BaggageBridging.toApplication(agentBaggage); - @SuppressWarnings("unchecked") - V value = (V) applicationBaggage; - return value; } + return applicationContext.get(key); } @Override public Context with(ContextKey k1, V v1) { - if (k1 == APPLICATION_SPAN_CONTEXT_KEY) { - Span applicationSpan = (Span) v1; - io.opentelemetry.api.trace.Span agentSpan = Bridging.toAgentOrNull(applicationSpan); - if (agentSpan == null) { - return this; + for (ContextKeyBridge bridge : CONTEXT_KEY_BRIDGES) { + Context context = bridge.with(this, k1, v1); + if (context != null) { + return context; } - return new AgentContextWrapper( - agentContext.with(AGENT_SPAN_CONTEXT_KEY, agentSpan), applicationContext); - } - if (k1 == APPLICATION_BAGGAGE_CONTEXT_KEY) { - Baggage applicationBaggage = (Baggage) v1; - io.opentelemetry.api.baggage.Baggage agentBaggage = - BaggageBridging.toAgent(applicationBaggage); - return new AgentContextWrapper( - agentContext.with(AGENT_BAGGAGE_CONTEXT_KEY, agentBaggage), applicationContext); } return new AgentContextWrapper(agentContext, applicationContext.with(k1, v1)); } @@ -230,4 +160,76 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable { return "{agentContext=" + agentContext + ", applicationContext=" + applicationContext + "}"; } } + + static class ContextKeyBridge { + + private final ContextKey applicationContextKey; + private final io.opentelemetry.context.ContextKey agentContextKey; + private final Function toAgent; + private final Function toApplication; + + @SuppressWarnings("unchecked") + ContextKeyBridge( + String applicationKeyHolderClassName, + String agentKeyHolderClassName, + Function toApplication, + Function toAgent) { + this.toApplication = toApplication; + this.toAgent = toAgent; + + ContextKey applicationContextKey; + try { + Class applicationKeyHolderClass = Class.forName(applicationKeyHolderClassName); + Field applicationContextKeyField = applicationKeyHolderClass.getDeclaredField("KEY"); + applicationContextKeyField.setAccessible(true); + applicationContextKey = (ContextKey) applicationContextKeyField.get(null); + } catch (Throwable t) { + applicationContextKey = null; + } + this.applicationContextKey = applicationContextKey; + + io.opentelemetry.context.ContextKey agentContextKey; + try { + Class agentKeyHolderClass = Class.forName(agentKeyHolderClassName); + Field agentContextKeyField = agentKeyHolderClass.getDeclaredField("KEY"); + agentContextKeyField.setAccessible(true); + agentContextKey = + (io.opentelemetry.context.ContextKey) agentContextKeyField.get(null); + } catch (Throwable t) { + agentContextKey = null; + } + this.agentContextKey = agentContextKey; + } + + @Nullable + V get(AgentContextWrapper contextWrapper, ContextKey requestedKey) { + if (requestedKey == applicationContextKey) { + AGENT agentValue = contextWrapper.agentContext.get(agentContextKey); + if (agentValue == null) { + return null; + } + APPLICATION applicationValue = toApplication.apply(agentValue); + @SuppressWarnings("unchecked") + V castValue = (V) applicationValue; + return castValue; + } + return null; + } + + @Nullable + Context with(AgentContextWrapper contextWrapper, ContextKey requestedKey, V value) { + if (requestedKey == applicationContextKey) { + @SuppressWarnings("unchecked") + APPLICATION applicationValue = (APPLICATION) value; + AGENT agentValue = toAgent.apply(applicationValue); + if (agentValue == null) { + return contextWrapper; + } + return new AgentContextWrapper( + contextWrapper.agentContext.with(agentContextKey, agentValue), + contextWrapper.applicationContext); + } + return null; + } + } }