Skip to content

Commit 3f3bb0f

Browse files
mccullsmcculls
committed
Support using the OpenTelemetry API to interact with automatic W3C baggage
Note: custom BaggageEntryMetadata is not supported
1 parent d0353e2 commit 3f3bb0f

File tree

7 files changed

+219
-0
lines changed

7 files changed

+219
-0
lines changed

‎dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/baggage/OtelBaggage.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
package datadog.opentelemetry.shim.baggage;
2+
3+
import static java.util.stream.Collectors.toMap;
4+
5+
import io.opentelemetry.api.baggage.Baggage;
6+
import io.opentelemetry.api.baggage.BaggageBuilder;
7+
import io.opentelemetry.api.baggage.BaggageEntry;
8+
import io.opentelemetry.api.baggage.BaggageEntryMetadata;
9+
import java.util.Map;
10+
import java.util.function.BiConsumer;
11+
import javax.annotation.Nullable;
12+
import javax.annotation.ParametersAreNonnullByDefault;
13+
14+
@ParametersAreNonnullByDefault
15+
public class OtelBaggage implements Baggage {
16+
private final datadog.trace.bootstrap.instrumentation.api.Baggage delegate;
17+
18+
public OtelBaggage(datadog.trace.bootstrap.instrumentation.api.Baggage delegate) {
19+
this.delegate = delegate;
20+
}
21+
22+
public OtelBaggage(Map<String, String> items) {
23+
this(datadog.trace.bootstrap.instrumentation.api.Baggage.create(items));
24+
}
25+
26+
@Override
27+
public int size() {
28+
return delegate.asMap().size();
29+
}
30+
31+
@Override
32+
public void forEach(BiConsumer<? super String, ? super BaggageEntry> consumer) {
33+
for (Map.Entry<String, String> entry : delegate.asMap().entrySet()) {
34+
consumer.accept(entry.getKey(), new ValueOnly(entry));
35+
}
36+
}
37+
38+
@Override
39+
public Map<String, BaggageEntry> asMap() {
40+
return delegate.asMap().entrySet().stream().collect(toMap(Map.Entry::getKey, ValueOnly::new));
41+
}
42+
43+
@Nullable
44+
@Override
45+
public String getEntryValue(String key) {
46+
return delegate.asMap().get(key);
47+
}
48+
49+
@Nullable
50+
@Override
51+
public BaggageEntry getEntry(String key) {
52+
String value = getEntryValue(key);
53+
return value != null ? new ValueOnly(value) : null;
54+
}
55+
56+
@Override
57+
public BaggageBuilder toBuilder() {
58+
return new OtelBaggageBuilder(delegate.asMap());
59+
}
60+
61+
public datadog.trace.bootstrap.instrumentation.api.Baggage asAgentBaggage() {
62+
return delegate;
63+
}
64+
65+
static class ValueOnly implements BaggageEntry {
66+
private final String value;
67+
68+
ValueOnly(String value) {
69+
this.value = value;
70+
}
71+
72+
ValueOnly(Map.Entry<String, String> entry) {
73+
this(entry.getValue());
74+
}
75+
76+
@Override
77+
public String getValue() {
78+
return value;
79+
}
80+
81+
@Override
82+
public BaggageEntryMetadata getMetadata() {
83+
return BaggageEntryMetadata.empty();
84+
}
85+
86+
@Override
87+
public int hashCode() {
88+
return value.hashCode();
89+
}
90+
91+
@Override
92+
public final boolean equals(Object o) {
93+
return (o instanceof ValueOnly) && value.equals(((ValueOnly) o).value);
94+
}
95+
}
96+
}

‎dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/baggage/OtelBaggageBuilder.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
package datadog.opentelemetry.shim.baggage;
2+
3+
import io.opentelemetry.api.baggage.Baggage;
4+
import io.opentelemetry.api.baggage.BaggageBuilder;
5+
import io.opentelemetry.api.baggage.BaggageEntryMetadata;
6+
import java.util.HashMap;
7+
import java.util.Map;
8+
import javax.annotation.Nullable;
9+
import javax.annotation.ParametersAreNonnullByDefault;
10+
11+
@ParametersAreNonnullByDefault
12+
public class OtelBaggageBuilder implements BaggageBuilder {
13+
private final Map<String, String> items;
14+
15+
public OtelBaggageBuilder(Map<String, String> items) {
16+
this.items = new HashMap<>(items);
17+
}
18+
19+
@Override
20+
public BaggageBuilder put(
21+
@Nullable String key, @Nullable String value, BaggageEntryMetadata ignore) {
22+
if (key != null && value != null) {
23+
items.put(key, value);
24+
}
25+
return this;
26+
}
27+
28+
@Override
29+
public BaggageBuilder remove(@Nullable String key) {
30+
if (key != null) {
31+
items.remove(key);
32+
}
33+
return this;
34+
}
35+
36+
@Override
37+
public Baggage build() {
38+
return new OtelBaggage(items);
39+
}
40+
}

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
package datadog.opentelemetry.shim.context;
22

3+
import datadog.opentelemetry.shim.baggage.OtelBaggage;
34
import datadog.opentelemetry.shim.trace.OtelSpan;
45
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
56
import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper;
7+
import datadog.trace.bootstrap.instrumentation.api.Baggage;
68
import io.opentelemetry.context.Context;
79
import io.opentelemetry.context.ContextKey;
810
import io.opentelemetry.context.Scope;
@@ -18,6 +20,7 @@ public class OtelContext implements Context {
1820
/** Overridden root context. */
1921
public static final Context ROOT = new OtelContext(datadog.context.Context.root());
2022

23+
private static final String OTEL_CONTEXT_BAGGAGE_KEY = "opentelemetry-baggage-key";
2124
private static final String OTEL_CONTEXT_SPAN_KEY = "opentelemetry-trace-span-key";
2225
private static final String OTEL_CONTEXT_ROOT_SPAN_KEY = "opentelemetry-traces-local-root-span";
2326

@@ -55,6 +58,12 @@ public <V> V get(ContextKey<V> key) {
5558
return (V) toOtelSpan(span.getLocalRootSpan());
5659
}
5760
// fall-through and check for non-datadog span data
61+
} else if (OTEL_CONTEXT_BAGGAGE_KEY.equals(key.toString())) {
62+
Baggage baggage = Baggage.fromContext(delegate);
63+
if (baggage != null) {
64+
return (V) new OtelBaggage(baggage);
65+
}
66+
// fall-through and check for non-datadog baggage
5867
}
5968
return (V) delegate.get(delegateKey(key));
6069
}
@@ -67,6 +76,12 @@ public <V> Context with(ContextKey<V> key, V value) {
6776
return new OtelContext(delegate.with(span));
6877
}
6978
// fall-through and store as non-datadog span data
79+
} else if (OTEL_CONTEXT_BAGGAGE_KEY.equals(key.toString())) {
80+
if (value instanceof OtelBaggage) {
81+
Baggage baggage = ((OtelBaggage) value).asAgentBaggage();
82+
return new OtelContext(delegate.with(baggage));
83+
}
84+
// fall-through and store as non-datadog baggage
7085
}
7186
return new OtelContext(delegate.with(delegateKey(key), value));
7287
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,9 @@ public String[] helperClassNames() {
6262
"datadog.opentelemetry.shim.context.propagation.AgentTextMapPropagator",
6363
"datadog.opentelemetry.shim.context.propagation.OtelContextPropagators",
6464
"datadog.opentelemetry.shim.context.propagation.TraceStateHelper",
65+
"datadog.opentelemetry.shim.baggage.OtelBaggage",
66+
"datadog.opentelemetry.shim.baggage.OtelBaggage$ValueOnly",
67+
"datadog.opentelemetry.shim.baggage.OtelBaggageBuilder",
6568
"datadog.opentelemetry.shim.trace.OtelExtractedContext",
6669
"datadog.opentelemetry.shim.trace.OtelConventions",
6770
"datadog.opentelemetry.shim.trace.OtelConventions$1",

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ public String[] helperClassNames() {
5656
return new String[] {
5757
"datadog.opentelemetry.shim.context.OtelContext",
5858
"datadog.opentelemetry.shim.context.OtelScope",
59+
"datadog.opentelemetry.shim.baggage.OtelBaggage",
60+
"datadog.opentelemetry.shim.baggage.OtelBaggage$ValueOnly",
61+
"datadog.opentelemetry.shim.baggage.OtelBaggageBuilder",
5962
"datadog.opentelemetry.shim.trace.OtelExtractedContext",
6063
"datadog.opentelemetry.shim.trace.OtelConventions",
6164
"datadog.opentelemetry.shim.trace.OtelConventions$1",

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,9 @@ public String[] helperClassNames() {
5757
return new String[] {
5858
"datadog.opentelemetry.shim.context.OtelContext",
5959
"datadog.opentelemetry.shim.context.OtelScope",
60+
"datadog.opentelemetry.shim.baggage.OtelBaggage",
61+
"datadog.opentelemetry.shim.baggage.OtelBaggage$ValueOnly",
62+
"datadog.opentelemetry.shim.baggage.OtelBaggageBuilder",
6063
"datadog.opentelemetry.shim.trace.OtelExtractedContext",
6164
"datadog.opentelemetry.shim.trace.OtelConventions",
6265
"datadog.opentelemetry.shim.trace.OtelConventions$1",

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

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package opentelemetry14.context
33
import datadog.trace.agent.test.InstrumentationSpecification
44
import datadog.trace.api.DDSpanId
55
import io.opentelemetry.api.GlobalOpenTelemetry
6+
import io.opentelemetry.api.baggage.Baggage
67
import io.opentelemetry.api.trace.Span
78
import io.opentelemetry.context.Context
89
import io.opentelemetry.context.ContextKey
@@ -324,6 +325,64 @@ class ContextTest extends InstrumentationSpecification {
324325
context.get(CustomData.KEY) == null
325326
}
326327

328+
def "test Baggage.current/makeCurrent()"() {
329+
when:
330+
def otelBaggage = Baggage.current()
331+
def otelBaggageFromContext = Baggage.fromContext(Context.current())
332+
def otelBaggageFromContextOrNull = Baggage.fromContextOrNull(Context.current())
333+
334+
then: "current baggage must be empty or null"
335+
otelBaggage != null
336+
otelBaggage.isEmpty()
337+
otelBaggageFromContext != null
338+
otelBaggageFromContext.isEmpty()
339+
otelBaggageFromContextOrNull == null
340+
341+
when:
342+
def ddContext = datadog.context.Context.current()
343+
def ddBaggage = datadog.trace.bootstrap.instrumentation.api.Baggage.create([
344+
"foo" : "value_to_be_replaced",
345+
"FOO" : "UNTOUCHED",
346+
"remove_me_key" : "remove_me_value"
347+
])
348+
def ddScope = ddContext.with(ddBaggage).attach()
349+
otelBaggage = Baggage.current()
350+
otelBaggageFromContext = Baggage.fromContext(Context.current())
351+
otelBaggageFromContextOrNull = Baggage.fromContextOrNull(Context.current())
352+
353+
then: "Datadog baggage must be current"
354+
otelBaggage != null
355+
otelBaggage.size() == 3
356+
otelBaggage.getEntryValue("foo") == "value_to_be_replaced"
357+
otelBaggage.getEntryValue("FOO") == "UNTOUCHED"
358+
otelBaggage.getEntryValue("remove_me_key") == "remove_me_value"
359+
otelBaggage.asMap() == otelBaggageFromContext.asMap()
360+
otelBaggage.asMap() == otelBaggageFromContextOrNull.asMap()
361+
362+
when:
363+
def builder = otelBaggage.toBuilder()
364+
builder.put("new_foo", "new_value")
365+
builder.put("foo", "overwrite_value")
366+
builder.remove("remove_me_key")
367+
def otelScope = builder.build().makeCurrent()
368+
otelBaggage = Baggage.current()
369+
otelBaggageFromContext = Baggage.fromContext(Context.current())
370+
otelBaggageFromContextOrNull = Baggage.fromContextOrNull(Context.current())
371+
372+
then: "baggage must contain OTel changes"
373+
otelBaggage != null
374+
otelBaggage.size() == 3
375+
otelBaggage.getEntryValue("foo") == "overwrite_value"
376+
otelBaggage.getEntryValue("FOO") == "UNTOUCHED"
377+
otelBaggage.getEntryValue("new_foo") == "new_value"
378+
otelBaggage.asMap() == otelBaggageFromContext.asMap()
379+
otelBaggage.asMap() == otelBaggageFromContextOrNull.asMap()
380+
381+
cleanup:
382+
otelScope.close()
383+
ddScope.close()
384+
}
385+
327386
@Override
328387
void cleanup() {
329388
// Test for context leak

0 commit comments

Comments
 (0)