diff --git a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java index 3a487ec30c9..bd3fc0ff17d 100644 --- a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java +++ b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java @@ -1,170 +1,109 @@ package datadog.opentelemetry.shim.context; import datadog.opentelemetry.shim.trace.OtelSpan; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper; -import io.opentelemetry.api.trace.Span; import io.opentelemetry.context.Context; import io.opentelemetry.context.ContextKey; import io.opentelemetry.context.Scope; -import java.util.Arrays; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import javax.annotation.Nullable; import javax.annotation.ParametersAreNonnullByDefault; +@SuppressWarnings({"rawtypes", "unchecked"}) @ParametersAreNonnullByDefault public class OtelContext implements Context { - private static final Object[] NO_ENTRIES = {}; /** Overridden root context. */ - public static final Context ROOT = new OtelContext(OtelSpan.invalid(), OtelSpan.invalid()); + public static final Context ROOT = new OtelContext(datadog.context.Context.root()); private static final String OTEL_CONTEXT_SPAN_KEY = "opentelemetry-trace-span-key"; private static final String OTEL_CONTEXT_ROOT_SPAN_KEY = "opentelemetry-traces-local-root-span"; - /** Keep track of propagated context that has not been captured on the scope stack. */ - private static final ThreadLocal lastPropagated = new ThreadLocal<>(); + /** Records the keys needed to access the delegate context, mapped by key name. */ + private static final Map, datadog.context.ContextKey> DELEGATE_KEYS = + new ConcurrentHashMap<>(); - private final Span currentSpan; - private final Span rootSpan; + private final datadog.context.Context delegate; - private final Object[] entries; + public OtelContext(datadog.context.Context delegate) { + this.delegate = delegate; + } - public OtelContext(Span currentSpan, Span rootSpan) { - this(currentSpan, rootSpan, NO_ENTRIES); + public static Context current() { + return new OtelContext(datadog.context.Context.current()); } - public OtelContext(Span currentSpan, Span rootSpan, Object[] entries) { - this.currentSpan = currentSpan; - this.rootSpan = rootSpan; - this.entries = entries; + @Override + public Scope makeCurrent() { + return new OtelScope(Context.super.makeCurrent(), delegate.attach()); } @Nullable @Override - @SuppressWarnings("unchecked") public V get(ContextKey key) { if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) { - return (V) this.currentSpan; + AgentSpan span = AgentSpan.fromContext(delegate); + if (span != null) { + return (V) toOtelSpan(span); + } + // fall-through and check for non-datadog span data } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) { - return (V) this.rootSpan; - } - for (int i = 0; i < this.entries.length; i += 2) { - if (this.entries[i] == key) { - return (V) this.entries[i + 1]; + AgentSpan span = AgentSpan.fromContext(delegate); + if (span != null) { + return (V) toOtelSpan(span.getLocalRootSpan()); } + // fall-through and check for non-datadog span data } - return null; + return (V) delegate.get(delegateKey(key)); } @Override public Context with(ContextKey key, V value) { if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) { - return new OtelContext((Span) value, this.rootSpan, this.entries); - } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) { - return new OtelContext(this.currentSpan, (Span) value, this.entries); - } - Object[] newEntries = null; - int oldEntriesLength = this.entries.length; - for (int i = 0; i < oldEntriesLength; i += 2) { - if (this.entries[i] == key) { - if (this.entries[i + 1] == value) { - return this; - } - newEntries = this.entries.clone(); - newEntries[i + 1] = value; - break; + if (value instanceof OtelSpan) { + AgentSpan span = ((OtelSpan) value).asAgentSpan(); + return new OtelContext(delegate.with(span)); } + // fall-through and store as non-datadog span data } - if (null == newEntries) { - newEntries = Arrays.copyOf(this.entries, oldEntriesLength + 2); - newEntries[oldEntriesLength] = key; - newEntries[oldEntriesLength + 1] = value; - } - return new OtelContext(this.currentSpan, this.rootSpan, newEntries); + return new OtelContext(delegate.with(delegateKey(key), value)); } @Override - public Scope makeCurrent() { - final Scope scope = Context.super.makeCurrent(); - if (this.currentSpan instanceof OtelSpan) { - // only keep propagated context until next span activation - lastPropagated.remove(); - AgentScope agentScope = ((OtelSpan) this.currentSpan).activate(); - return new OtelScope(scope, agentScope, this.entries); - } else { - // propagated context not on the scope stack, capture it here - lastPropagated.set(this); - return new Scope() { - @Override - public void close() { - lastPropagated.remove(); - scope.close(); - } - }; - } - } - - public static Context current() { - // Check for propagated context not on the scope stack - Context context = lastPropagated.get(); - if (null != context) { - return context; - } - // Check empty context - AgentScope agentCurrentScope = AgentTracer.activeScope(); - if (null == agentCurrentScope) { - return OtelContext.ROOT; - } - // Get OTel current span - Span otelCurrentSpan = null; - AgentSpan agentCurrentSpan = agentCurrentScope.span(); - if (agentCurrentSpan instanceof AttachableWrapper) { - Object wrapper = ((AttachableWrapper) agentCurrentSpan).getWrapper(); - if (wrapper instanceof OtelSpan) { - otelCurrentSpan = (OtelSpan) wrapper; - } + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) { + return false; } - if (otelCurrentSpan == null) { - otelCurrentSpan = new OtelSpan(agentCurrentSpan); - } - // Get OTel root span - Span otelRootSpan = null; - AgentSpan agentRootSpan = agentCurrentSpan.getLocalRootSpan(); - if (agentRootSpan instanceof AttachableWrapper) { - Object wrapper = ((AttachableWrapper) agentRootSpan).getWrapper(); - if (wrapper instanceof OtelSpan) { - otelRootSpan = (OtelSpan) wrapper; - } - } - if (otelRootSpan == null) { - otelRootSpan = new OtelSpan(agentRootSpan); - } - // Get OTel custom context entries - Object[] contextEntries = NO_ENTRIES; - if (agentCurrentScope instanceof AttachableWrapper) { - Object wrapper = ((AttachableWrapper) agentCurrentScope).getWrapper(); - if (wrapper instanceof OtelScope) { - contextEntries = ((OtelScope) wrapper).contextEntries(); - } - } - return new OtelContext(otelCurrentSpan, otelRootSpan, contextEntries); + return delegate.equals(((OtelContext) o).delegate); } - /** Last propagated context not on the scope stack; {@code null} if there's no such context. */ - @Nullable - public static Context lastPropagated() { - return lastPropagated.get(); + @Override + public int hashCode() { + return delegate.hashCode(); } @Override public String toString() { - return "OtelContext{" - + "currentSpan=" - + this.currentSpan.getSpanContext() - + ", rootSpan=" - + this.rootSpan.getSpanContext() - + '}'; + return "OtelContext{" + "delegate=" + delegate + '}'; + } + + private static datadog.context.ContextKey delegateKey(ContextKey key) { + return DELEGATE_KEYS.computeIfAbsent(key, OtelContext::mapByKeyName); + } + + private static datadog.context.ContextKey mapByKeyName(ContextKey key) { + return datadog.context.ContextKey.named(key.toString()); + } + + private static OtelSpan toOtelSpan(AgentSpan span) { + if (span instanceof AttachableWrapper) { + Object wrapper = ((AttachableWrapper) span).getWrapper(); + if (wrapper instanceof OtelSpan) { + return (OtelSpan) wrapper; + } + } + return new OtelSpan(span); } } diff --git a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelScope.java b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelScope.java index f5570b83c6b..db28fa46606 100644 --- a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelScope.java +++ b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelScope.java @@ -1,31 +1,20 @@ package datadog.opentelemetry.shim.context; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; -import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper; +import datadog.context.ContextScope; import io.opentelemetry.context.Scope; public class OtelScope implements Scope { private final Scope scope; - private final AgentScope delegate; - private final Object[] contextEntries; + private final ContextScope delegate; - public OtelScope(Scope scope, AgentScope delegate, Object[] contextEntries) { + public OtelScope(Scope scope, ContextScope delegate) { this.scope = scope; this.delegate = delegate; - this.contextEntries = contextEntries; - if (delegate instanceof AttachableWrapper) { - ((AttachableWrapper) delegate).attachWrapper(this); - } - } - - /** Context entries from {@link OtelContext}, captured when the context was made current. */ - Object[] contextEntries() { - return contextEntries; } @Override public void close() { - this.delegate.close(); - this.scope.close(); + delegate.close(); + scope.close(); } } diff --git a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java index 86b6bf7986f..71feb6b9f56 100644 --- a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java +++ b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java @@ -7,7 +7,6 @@ import datadog.opentelemetry.shim.context.OtelContext; import datadog.opentelemetry.shim.trace.OtelExtractedContext; -import datadog.opentelemetry.shim.trace.OtelSpan; import datadog.trace.api.TracePropagationStyle; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; @@ -59,7 +58,7 @@ public Context extract(Context context, @Nullable C carrier, TextMapGetter