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 a9d0064b75e..28d1308c7f7 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 @@ -9,11 +9,14 @@ import io.opentelemetry.context.Context; import io.opentelemetry.context.ContextKey; import io.opentelemetry.context.Scope; +import java.util.Arrays; import javax.annotation.Nullable; import javax.annotation.ParametersAreNonnullByDefault; @ParametersAreNonnullByDefault public class OtelContext implements Context { + private static final Object[] NO_ENTRIES = {}; + /** Overridden root context. */ public static final OtelContext ROOT = new OtelContext(OtelSpan.invalid(), OtelSpan.invalid()); @@ -26,30 +29,60 @@ public class OtelContext implements Context { private final Span currentSpan; private final Span rootSpan; + private final Object[] entries; + public OtelContext(Span currentSpan, Span rootSpan) { + this(currentSpan, rootSpan, NO_ENTRIES); + } + + public OtelContext(Span currentSpan, Span rootSpan, Object[] entries) { this.currentSpan = currentSpan; this.rootSpan = rootSpan; + this.entries = entries; } @Nullable @Override + @SuppressWarnings("unchecked") public V get(ContextKey key) { if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) { return (V) this.currentSpan; } 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]; + } + } return null; } @Override - public Context with(ContextKey k1, V v1) { - if (OTEL_CONTEXT_SPAN_KEY.equals(k1.toString())) { - return new OtelContext((Span) v1, this.rootSpan); - } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(k1.toString())) { - return new OtelContext(this.currentSpan, (Span) v1); + 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; + } } - return this; + 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); } @Override @@ -59,7 +92,7 @@ public Scope makeCurrent() { // only keep propagated context until next span activation lastPropagated.remove(); AgentScope agentScope = ((OtelSpan) this.currentSpan).activate(); - return new OtelScope(scope, agentScope); + return new OtelScope(scope, agentScope, this.entries); } else { // propagated context not on the scope stack, capture it here lastPropagated.set(this); @@ -80,12 +113,13 @@ public static Context current() { return context; } // Check empty context - AgentSpan agentCurrentSpan = AgentTracer.activeSpan(); - if (null == agentCurrentSpan) { + 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) { @@ -107,7 +141,15 @@ public static Context current() { if (otelRootSpan == null) { otelRootSpan = new OtelSpan(agentRootSpan); } - return new OtelContext(otelCurrentSpan, otelRootSpan); + // 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); } /** Last propagated context not on the scope stack; {@code null} if there's no such context. */ 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 7ff94ca0f77..f5570b83c6b 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 @@ -7,15 +7,22 @@ public class OtelScope implements Scope { private final Scope scope; private final AgentScope delegate; + private final Object[] contextEntries; - public OtelScope(Scope scope, AgentScope delegate) { + public OtelScope(Scope scope, AgentScope delegate, Object[] contextEntries) { 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(); diff --git a/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextTest.groovy b/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextTest.groovy index ab386e4b7e2..023cdce1643 100644 --- a/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextTest.groovy +++ b/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextTest.groovy @@ -8,7 +8,6 @@ import io.opentelemetry.context.Context import io.opentelemetry.context.ContextKey import io.opentelemetry.context.ImplicitContextKeyed import io.opentelemetry.context.ThreadLocalContextStorage -import spock.lang.Ignore import spock.lang.Subject import static datadog.trace.bootstrap.instrumentation.api.ScopeSource.MANUAL @@ -279,7 +278,6 @@ class ContextTest extends AgentTestRunner { parentSpan.end() } - @Ignore("Not supported") def "test custom object storage"() { setup: def context = Context.root()