Skip to content

Commit d8de104

Browse files
mccullsmcculls
authored
Migrate OtelContext wrapper to new Context API (#8645)
1 parent e8f5aae commit d8de104

File tree

8 files changed

+77
-146
lines changed

8 files changed

+77
-146
lines changed

‎dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java‎

Lines changed: 57 additions & 118 deletions
Original file line numberDiff line numberDiff line change
@@ -1,170 +1,109 @@
11
package datadog.opentelemetry.shim.context;
22

33
import datadog.opentelemetry.shim.trace.OtelSpan;
4-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
54
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
6-
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
75
import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper;
8-
import io.opentelemetry.api.trace.Span;
96
import io.opentelemetry.context.Context;
107
import io.opentelemetry.context.ContextKey;
118
import io.opentelemetry.context.Scope;
12-
import java.util.Arrays;
9+
import java.util.Map;
10+
import java.util.concurrent.ConcurrentHashMap;
1311
import javax.annotation.Nullable;
1412
import javax.annotation.ParametersAreNonnullByDefault;
1513

14+
@SuppressWarnings({"rawtypes", "unchecked"})
1615
@ParametersAreNonnullByDefault
1716
public class OtelContext implements Context {
18-
private static final Object[] NO_ENTRIES = {};
1917

2018
/** Overridden root context. */
21-
public static final Context ROOT = new OtelContext(OtelSpan.invalid(), OtelSpan.invalid());
19+
public static final Context ROOT = new OtelContext(datadog.context.Context.root());
2220

2321
private static final String OTEL_CONTEXT_SPAN_KEY = "opentelemetry-trace-span-key";
2422
private static final String OTEL_CONTEXT_ROOT_SPAN_KEY = "opentelemetry-traces-local-root-span";
2523

26-
/** Keep track of propagated context that has not been captured on the scope stack. */
27-
private static final ThreadLocal<OtelContext> lastPropagated = new ThreadLocal<>();
24+
/** Records the keys needed to access the delegate context, mapped by key name. */
25+
private static final Map<ContextKey<?>, datadog.context.ContextKey<?>> DELEGATE_KEYS =
26+
new ConcurrentHashMap<>();
2827

29-
private final Span currentSpan;
30-
private final Span rootSpan;
28+
private final datadog.context.Context delegate;
3129

32-
private final Object[] entries;
30+
public OtelContext(datadog.context.Context delegate) {
31+
this.delegate = delegate;
32+
}
3333

34-
public OtelContext(Span currentSpan, Span rootSpan) {
35-
this(currentSpan, rootSpan, NO_ENTRIES);
34+
public static Context current() {
35+
return new OtelContext(datadog.context.Context.current());
3636
}
3737

38-
public OtelContext(Span currentSpan, Span rootSpan, Object[] entries) {
39-
this.currentSpan = currentSpan;
40-
this.rootSpan = rootSpan;
41-
this.entries = entries;
38+
@Override
39+
public Scope makeCurrent() {
40+
return new OtelScope(Context.super.makeCurrent(), delegate.attach());
4241
}
4342

4443
@Nullable
4544
@Override
46-
@SuppressWarnings("unchecked")
4745
public <V> V get(ContextKey<V> key) {
4846
if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) {
49-
return (V) this.currentSpan;
47+
AgentSpan span = AgentSpan.fromContext(delegate);
48+
if (span != null) {
49+
return (V) toOtelSpan(span);
50+
}
51+
// fall-through and check for non-datadog span data
5052
} else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) {
51-
return (V) this.rootSpan;
52-
}
53-
for (int i = 0; i < this.entries.length; i += 2) {
54-
if (this.entries[i] == key) {
55-
return (V) this.entries[i + 1];
53+
AgentSpan span = AgentSpan.fromContext(delegate);
54+
if (span != null) {
55+
return (V) toOtelSpan(span.getLocalRootSpan());
5656
}
57+
// fall-through and check for non-datadog span data
5758
}
58-
return null;
59+
return (V) delegate.get(delegateKey(key));
5960
}
6061

6162
@Override
6263
public <V> Context with(ContextKey<V> key, V value) {
6364
if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) {
64-
return new OtelContext((Span) value, this.rootSpan, this.entries);
65-
} else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) {
66-
return new OtelContext(this.currentSpan, (Span) value, this.entries);
67-
}
68-
Object[] newEntries = null;
69-
int oldEntriesLength = this.entries.length;
70-
for (int i = 0; i < oldEntriesLength; i += 2) {
71-
if (this.entries[i] == key) {
72-
if (this.entries[i + 1] == value) {
73-
return this;
74-
}
75-
newEntries = this.entries.clone();
76-
newEntries[i + 1] = value;
77-
break;
65+
if (value instanceof OtelSpan) {
66+
AgentSpan span = ((OtelSpan) value).asAgentSpan();
67+
return new OtelContext(delegate.with(span));
7868
}
69+
// fall-through and store as non-datadog span data
7970
}
80-
if (null == newEntries) {
81-
newEntries = Arrays.copyOf(this.entries, oldEntriesLength + 2);
82-
newEntries[oldEntriesLength] = key;
83-
newEntries[oldEntriesLength + 1] = value;
84-
}
85-
return new OtelContext(this.currentSpan, this.rootSpan, newEntries);
71+
return new OtelContext(delegate.with(delegateKey(key), value));
8672
}
8773

8874
@Override
89-
public Scope makeCurrent() {
90-
final Scope scope = Context.super.makeCurrent();
91-
if (this.currentSpan instanceof OtelSpan) {
92-
// only keep propagated context until next span activation
93-
lastPropagated.remove();
94-
AgentScope agentScope = ((OtelSpan) this.currentSpan).activate();
95-
return new OtelScope(scope, agentScope, this.entries);
96-
} else {
97-
// propagated context not on the scope stack, capture it here
98-
lastPropagated.set(this);
99-
return new Scope() {
100-
@Override
101-
public void close() {
102-
lastPropagated.remove();
103-
scope.close();
104-
}
105-
};
106-
}
107-
}
108-
109-
public static Context current() {
110-
// Check for propagated context not on the scope stack
111-
Context context = lastPropagated.get();
112-
if (null != context) {
113-
return context;
114-
}
115-
// Check empty context
116-
AgentScope agentCurrentScope = AgentTracer.activeScope();
117-
if (null == agentCurrentScope) {
118-
return OtelContext.ROOT;
119-
}
120-
// Get OTel current span
121-
Span otelCurrentSpan = null;
122-
AgentSpan agentCurrentSpan = agentCurrentScope.span();
123-
if (agentCurrentSpan instanceof AttachableWrapper) {
124-
Object wrapper = ((AttachableWrapper) agentCurrentSpan).getWrapper();
125-
if (wrapper instanceof OtelSpan) {
126-
otelCurrentSpan = (OtelSpan) wrapper;
127-
}
75+
public boolean equals(Object o) {
76+
if (o == null || getClass() != o.getClass()) {
77+
return false;
12878
}
129-
if (otelCurrentSpan == null) {
130-
otelCurrentSpan = new OtelSpan(agentCurrentSpan);
131-
}
132-
// Get OTel root span
133-
Span otelRootSpan = null;
134-
AgentSpan agentRootSpan = agentCurrentSpan.getLocalRootSpan();
135-
if (agentRootSpan instanceof AttachableWrapper) {
136-
Object wrapper = ((AttachableWrapper) agentRootSpan).getWrapper();
137-
if (wrapper instanceof OtelSpan) {
138-
otelRootSpan = (OtelSpan) wrapper;
139-
}
140-
}
141-
if (otelRootSpan == null) {
142-
otelRootSpan = new OtelSpan(agentRootSpan);
143-
}
144-
// Get OTel custom context entries
145-
Object[] contextEntries = NO_ENTRIES;
146-
if (agentCurrentScope instanceof AttachableWrapper) {
147-
Object wrapper = ((AttachableWrapper) agentCurrentScope).getWrapper();
148-
if (wrapper instanceof OtelScope) {
149-
contextEntries = ((OtelScope) wrapper).contextEntries();
150-
}
151-
}
152-
return new OtelContext(otelCurrentSpan, otelRootSpan, contextEntries);
79+
return delegate.equals(((OtelContext) o).delegate);
15380
}
15481

155-
/** Last propagated context not on the scope stack; {@code null} if there's no such context. */
156-
@Nullable
157-
public static Context lastPropagated() {
158-
return lastPropagated.get();
82+
@Override
83+
public int hashCode() {
84+
return delegate.hashCode();
15985
}
16086

16187
@Override
16288
public String toString() {
163-
return "OtelContext{"
164-
+ "currentSpan="
165-
+ this.currentSpan.getSpanContext()
166-
+ ", rootSpan="
167-
+ this.rootSpan.getSpanContext()
168-
+ '}';
89+
return "OtelContext{" + "delegate=" + delegate + '}';
90+
}
91+
92+
private static datadog.context.ContextKey delegateKey(ContextKey key) {
93+
return DELEGATE_KEYS.computeIfAbsent(key, OtelContext::mapByKeyName);
94+
}
95+
96+
private static datadog.context.ContextKey mapByKeyName(ContextKey key) {
97+
return datadog.context.ContextKey.named(key.toString());
98+
}
99+
100+
private static OtelSpan toOtelSpan(AgentSpan span) {
101+
if (span instanceof AttachableWrapper) {
102+
Object wrapper = ((AttachableWrapper) span).getWrapper();
103+
if (wrapper instanceof OtelSpan) {
104+
return (OtelSpan) wrapper;
105+
}
106+
}
107+
return new OtelSpan(span);
169108
}
170109
}

‎dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelScope.java‎

Lines changed: 5 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,20 @@
11
package datadog.opentelemetry.shim.context;
22

3-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
4-
import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper;
3+
import datadog.context.ContextScope;
54
import io.opentelemetry.context.Scope;
65

76
public class OtelScope implements Scope {
87
private final Scope scope;
9-
private final AgentScope delegate;
10-
private final Object[] contextEntries;
8+
private final ContextScope delegate;
119

12-
public OtelScope(Scope scope, AgentScope delegate, Object[] contextEntries) {
10+
public OtelScope(Scope scope, ContextScope delegate) {
1311
this.scope = scope;
1412
this.delegate = delegate;
15-
this.contextEntries = contextEntries;
16-
if (delegate instanceof AttachableWrapper) {
17-
((AttachableWrapper) delegate).attachWrapper(this);
18-
}
19-
}
20-
21-
/** Context entries from {@link OtelContext}, captured when the context was made current. */
22-
Object[] contextEntries() {
23-
return contextEntries;
2413
}
2514

2615
@Override
2716
public void close() {
28-
this.delegate.close();
29-
this.scope.close();
17+
delegate.close();
18+
scope.close();
3019
}
3120
}

‎dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77

88
import datadog.opentelemetry.shim.context.OtelContext;
99
import datadog.opentelemetry.shim.trace.OtelExtractedContext;
10-
import datadog.opentelemetry.shim.trace.OtelSpan;
1110
import datadog.trace.api.TracePropagationStyle;
1211
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1312
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
@@ -59,7 +58,7 @@ public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C
5958
} else {
6059
TraceState traceState = extractTraceState(extracted, carrier, getter);
6160
SpanContext spanContext = fromRemote(extracted, traceState);
62-
return new OtelContext(Span.wrap(spanContext), OtelSpan.invalid());
61+
return Span.wrap(spanContext).storeInContext(OtelContext.ROOT);
6362
}
6463
}
6564

‎dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/trace/OtelSpanBuilder.java‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
import static java.lang.Boolean.parseBoolean;
1212
import static java.util.Locale.ROOT;
1313

14-
import datadog.opentelemetry.shim.context.OtelContext;
1514
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1615
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
1716
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -202,11 +201,7 @@ public Span startSpan() {
202201
setSpanKind(INTERNAL);
203202
}
204203
if (!this.ignoreActiveSpan) {
205-
// support automatic parenting from propagated context not on the scope stack
206-
Context context = OtelContext.lastPropagated();
207-
if (null != context) {
208-
setParent(context);
209-
}
204+
setParent(Context.current());
210205
}
211206
AgentSpan delegate = this.delegate.start();
212207
// Apply overrides

‎dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/main/java/datadog/trace/instrumentation/opentelemetry14/OpenTelemetryInstrumentation.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,6 @@ public boolean onlyMatchKnownTypes() {
5858
public String[] helperClassNames() {
5959
return new String[] {
6060
"datadog.opentelemetry.shim.context.OtelContext",
61-
"datadog.opentelemetry.shim.context.OtelContext$1",
6261
"datadog.opentelemetry.shim.context.OtelScope",
6362
"datadog.opentelemetry.shim.context.propagation.AgentTextMapPropagator",
6463
"datadog.opentelemetry.shim.context.propagation.OtelContextPropagators",

‎dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/main/java/datadog/trace/instrumentation/opentelemetry14/context/OpenTelemetryContextInstrumentation.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,6 @@ public boolean onlyMatchKnownTypes() {
5555
public String[] helperClassNames() {
5656
return new String[] {
5757
"datadog.opentelemetry.shim.context.OtelContext",
58-
"datadog.opentelemetry.shim.context.OtelContext$1",
5958
"datadog.opentelemetry.shim.context.OtelScope",
6059
"datadog.opentelemetry.shim.trace.OtelExtractedContext",
6160
"datadog.opentelemetry.shim.trace.OtelConventions",

‎dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/main/java/datadog/trace/instrumentation/opentelemetry14/context/OpenTelemetryContextStorageInstrumentation.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,6 @@ public boolean onlyMatchKnownTypes() {
5656
public String[] helperClassNames() {
5757
return new String[] {
5858
"datadog.opentelemetry.shim.context.OtelContext",
59-
"datadog.opentelemetry.shim.context.OtelContext$1",
6059
"datadog.opentelemetry.shim.context.OtelScope",
6160
"datadog.opentelemetry.shim.trace.OtelExtractedContext",
6261
"datadog.opentelemetry.shim.trace.OtelConventions",

‎dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextTest.groovy‎ renamed to ‎dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextForkedTest.groovy‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import static datadog.opentelemetry.shim.context.OtelContext.OTEL_CONTEXT_ROOT_S
1414
import static datadog.opentelemetry.shim.context.OtelContext.OTEL_CONTEXT_SPAN_KEY
1515
import static datadog.opentelemetry.shim.trace.OtelConventions.SPAN_KIND_INTERNAL
1616

17-
class ContextTest extends AgentTestRunner {
17+
class ContextForkedTest extends AgentTestRunner {
1818
@Subject
1919
def tracer = GlobalOpenTelemetry.get().tracerProvider.get("context-instrumentation")
2020

@@ -169,6 +169,10 @@ class ContextTest extends AgentTestRunner {
169169
then:
170170
currentSpan != null
171171
!currentSpan.spanContext.isValid()
172+
173+
cleanup:
174+
ddScope.close()
175+
ddSpan.finish()
172176
}
173177

174178
def "test clearing context"() {
@@ -238,6 +242,14 @@ class ContextTest extends AgentTestRunner {
238242
}
239243
}
240244
}
245+
246+
cleanup:
247+
otelGrandChildScope?.close()
248+
otelGrandChildSpan?.end()
249+
ddChildScope?.close()
250+
ddChildSpan?.finish()
251+
otelParentScope.close()
252+
otelParentSpan.end()
241253
}
242254

243255
def "test context spans retrieval"() {

0 commit comments

Comments
 (0)