|
1 | 1 | package datadog.opentelemetry.shim.context; |
2 | 2 |
|
3 | 3 | import datadog.opentelemetry.shim.trace.OtelSpan; |
4 | | -import datadog.trace.bootstrap.instrumentation.api.AgentScope; |
5 | 4 | import datadog.trace.bootstrap.instrumentation.api.AgentSpan; |
6 | | -import datadog.trace.bootstrap.instrumentation.api.AgentTracer; |
7 | 5 | import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper; |
8 | | -import io.opentelemetry.api.trace.Span; |
9 | 6 | import io.opentelemetry.context.Context; |
10 | 7 | import io.opentelemetry.context.ContextKey; |
11 | 8 | import io.opentelemetry.context.Scope; |
12 | | -import java.util.Arrays; |
| 9 | +import java.util.Map; |
| 10 | +import java.util.concurrent.ConcurrentHashMap; |
13 | 11 | import javax.annotation.Nullable; |
14 | 12 | import javax.annotation.ParametersAreNonnullByDefault; |
15 | 13 |
|
| 14 | +@SuppressWarnings({"rawtypes", "unchecked"}) |
16 | 15 | @ParametersAreNonnullByDefault |
17 | 16 | public class OtelContext implements Context { |
18 | | - private static final Object[] NO_ENTRIES = {}; |
19 | 17 |
|
20 | 18 | /** 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()); |
22 | 20 |
|
23 | 21 | private static final String OTEL_CONTEXT_SPAN_KEY = "opentelemetry-trace-span-key"; |
24 | 22 | private static final String OTEL_CONTEXT_ROOT_SPAN_KEY = "opentelemetry-traces-local-root-span"; |
25 | 23 |
|
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 | + private static final Map<ContextKey<?>, datadog.context.ContextKey<?>> CONTEXT_KEYS = |
| 25 | + new ConcurrentHashMap<>(); |
28 | 26 |
|
29 | | - private final Span currentSpan; |
30 | | - private final Span rootSpan; |
| 27 | + private final datadog.context.Context delegate; |
31 | 28 |
|
32 | | - private final Object[] entries; |
| 29 | + public OtelContext(datadog.context.Context delegate) { |
| 30 | + this.delegate = delegate; |
| 31 | + } |
33 | 32 |
|
34 | | - public OtelContext(Span currentSpan, Span rootSpan) { |
35 | | - this(currentSpan, rootSpan, NO_ENTRIES); |
| 33 | + public static Context current() { |
| 34 | + return new OtelContext(datadog.context.Context.current()); |
36 | 35 | } |
37 | 36 |
|
38 | | - public OtelContext(Span currentSpan, Span rootSpan, Object[] entries) { |
39 | | - this.currentSpan = currentSpan; |
40 | | - this.rootSpan = rootSpan; |
41 | | - this.entries = entries; |
| 37 | + @Override |
| 38 | + public Scope makeCurrent() { |
| 39 | + return new OtelScope(Context.super.makeCurrent(), delegate.attach()); |
42 | 40 | } |
43 | 41 |
|
44 | 42 | @Nullable |
45 | 43 | @Override |
46 | | - @SuppressWarnings("unchecked") |
47 | 44 | public <V> V get(ContextKey<V> key) { |
48 | 45 | if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) { |
49 | | - return (V) this.currentSpan; |
50 | | - } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) { |
51 | | - return (V) this.rootSpan; |
| 46 | + AgentSpan span = AgentSpan.fromContext(delegate); |
| 47 | + if (span != null) { |
| 48 | + return (V) toOtelSpan(span); |
| 49 | + } |
52 | 50 | } |
53 | | - for (int i = 0; i < this.entries.length; i += 2) { |
54 | | - if (this.entries[i] == key) { |
55 | | - return (V) this.entries[i + 1]; |
| 51 | + if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) { |
| 52 | + AgentSpan span = AgentSpan.fromContext(delegate); |
| 53 | + if (span != null) { |
| 54 | + return (V) toOtelSpan(span.getLocalRootSpan()); |
56 | 55 | } |
57 | 56 | } |
58 | | - return null; |
| 57 | + return (V) delegate.get(datadogKey(key)); |
59 | 58 | } |
60 | 59 |
|
61 | 60 | @Override |
62 | 61 | public <V> Context with(ContextKey<V> key, V value) { |
63 | 62 | 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; |
| 63 | + if (value instanceof OtelSpan) { |
| 64 | + AgentSpan span = ((OtelSpan) value).asAgentSpan(); |
| 65 | + return new OtelContext(delegate.with(span)); |
78 | 66 | } |
79 | 67 | } |
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); |
| 68 | + return new OtelContext(delegate.with(datadogKey(key), value)); |
86 | 69 | } |
87 | 70 |
|
88 | 71 | @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 | | - } |
128 | | - } |
129 | | - if (otelCurrentSpan == null) { |
130 | | - otelCurrentSpan = new OtelSpan(agentCurrentSpan); |
| 72 | + public boolean equals(Object o) { |
| 73 | + if (o == null || getClass() != o.getClass()) { |
| 74 | + return false; |
131 | 75 | } |
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); |
| 76 | + return delegate.equals(((OtelContext) o).delegate); |
153 | 77 | } |
154 | 78 |
|
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(); |
| 79 | + @Override |
| 80 | + public int hashCode() { |
| 81 | + return delegate.hashCode(); |
159 | 82 | } |
160 | 83 |
|
161 | 84 | @Override |
162 | 85 | public String toString() { |
163 | | - return "OtelContext{" |
164 | | - + "currentSpan=" |
165 | | - + this.currentSpan.getSpanContext() |
166 | | - + ", rootSpan=" |
167 | | - + this.rootSpan.getSpanContext() |
168 | | - + '}'; |
| 86 | + return "OtelContext{" + "delegate=" + delegate + '}'; |
| 87 | + } |
| 88 | + |
| 89 | + static datadog.context.ContextKey datadogKey(ContextKey key) { |
| 90 | + return CONTEXT_KEYS.computeIfAbsent(key, OtelContext::mapKeyByName); |
| 91 | + } |
F438
| 92 | + |
| 93 | + private static datadog.context.ContextKey mapKeyByName(ContextKey key) { |
| 94 | + return datadog.context.ContextKey.named(key.toString()); |
| 95 | + } |
| 96 | + |
| 97 | + private static OtelSpan toOtelSpan(AgentSpan span) { |
| 98 | + if (span instanceof AttachableWrapper) { |
| 99 | + Object wrapper = ((AttachableWrapper) span).getWrapper(); |
| 100 | + if (wrapper instanceof OtelSpan) { |
| 101 | + return (OtelSpan) wrapper; |
| 102 | + } |
| 103 | + } |
| 104 | + return new OtelSpan(span); |
169 | 105 | } |
170 | 106 | } |
0 commit comments