diff --git a/components/context/src/main/java/datadog/context/Context.java b/components/context/src/main/java/datadog/context/Context.java index add43186d44..da04ac4fe0b 100644 --- a/components/context/src/main/java/datadog/context/Context.java +++ b/components/context/src/main/java/datadog/context/Context.java @@ -114,7 +114,7 @@ static Context detachFrom(Object carrier) { /** * Creates a copy of this context with the given key-value set. * - *

Existing value with the given key will be replaced, and mapping to a {@code null} value will + *

Existing value with the given key will be replaced. Mapping to a {@code null} value will * remove the key-value from the context copy. * * @param the type of the value. @@ -124,6 +124,28 @@ static Context detachFrom(Object carrier) { */ Context with(ContextKey key, @Nullable T value); + /** + * Creates a copy of this context with the given pair of key-values. + * + *

Existing values with the given keys will be replaced. Mapping to a {@code null} value will + * remove the key-value from the context copy. + * + * @param the type of the first value. + * @param the type of the second value. + * @param firstKey the first key to store the first value. + * @param firstValue the first value to store. + * @param secondKey the second key to store the second value. + * @param secondValue the second value to store. + * @return a new context with the pair of key-values set. + */ + default Context with( + ContextKey firstKey, + @Nullable T firstValue, + ContextKey secondKey, + @Nullable U secondValue) { + return with(firstKey, firstValue).with(secondKey, secondValue); + } + /** * Creates a copy of this context with the implicit key is mapped to the value. * diff --git a/components/context/src/test/java/datadog/context/ContextTest.java b/components/context/src/test/java/datadog/context/ContextTest.java index 1fbb58643c0..52cc7da0371 100644 --- a/components/context/src/test/java/datadog/context/ContextTest.java +++ b/components/context/src/test/java/datadog/context/ContextTest.java @@ -63,6 +63,41 @@ void testWith(Context context) { assertDoesNotThrow(() -> context.with(null), "Null implicitly keyed value not throw exception"); } + @ParameterizedTest + @MethodSource("contextImplementations") + void testWithPair(Context context) { + // Test retrieving value + String stringValue = "value"; + Context context1 = context.with(BOOLEAN_KEY, false, STRING_KEY, stringValue); + assertEquals(stringValue, context1.get(STRING_KEY)); + assertEquals(false, context1.get(BOOLEAN_KEY)); + // Test overriding value + String stringValue2 = "value2"; + Context context2 = context1.with(STRING_KEY, stringValue2, BOOLEAN_KEY, true); + assertEquals(stringValue2, context2.get(STRING_KEY)); + assertEquals(true, context2.get(BOOLEAN_KEY)); + // Test clearing value + Context context3 = context2.with(BOOLEAN_KEY, null, STRING_KEY, null); + assertNull(context3.get(STRING_KEY)); + assertNull(context3.get(BOOLEAN_KEY)); + // Test null key handling + assertThrows( + NullPointerException.class, + () -> context.with(null, "test", STRING_KEY, "test"), + "Context forbids null keys"); + assertThrows( + NullPointerException.class, + () -> context.with(STRING_KEY, "test", null, "test"), + "Context forbids null keys"); + // Test null value handling + assertDoesNotThrow( + () -> context.with(BOOLEAN_KEY, null, STRING_KEY, "test"), + "Null value should not throw exception"); + assertDoesNotThrow( + () -> context.with(STRING_KEY, "test", BOOLEAN_KEY, null), + "Null value should not throw exception"); + } + @ParameterizedTest @MethodSource("contextImplementations") void testGet(Context original) { diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Constants.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Constants.java index 2255a531095..96be734ba22 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Constants.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Constants.java @@ -15,6 +15,7 @@ public final class Constants { */ public static final String[] BOOTSTRAP_PACKAGE_PREFIXES = { "datadog.slf4j", + "datadog.context", "datadog.appsec.api", "datadog.trace.api", "datadog.trace.bootstrap", diff --git a/dd-java-agent/build.gradle b/dd-java-agent/build.gradle index d8ce652266e..459880a0ee7 100644 --- a/dd-java-agent/build.gradle +++ b/dd-java-agent/build.gradle @@ -135,6 +135,7 @@ def sharedShadowJar = tasks.register('sharedShadowJar', ShadowJar) { exclude(project(':dd-java-agent:agent-logging')) exclude(project(':dd-trace-api')) exclude(project(':internal-api')) + exclude(project(':components:context')) exclude(project(':utils:time-utils')) exclude(dependency('org.slf4j::')) } diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/SpockRunner.java b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/SpockRunner.java index bb2f9ab0b3c..4fb7083f238 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/SpockRunner.java +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/SpockRunner.java @@ -38,6 +38,7 @@ public class SpockRunner extends JUnitPlatform { */ public static final String[] BOOTSTRAP_PACKAGE_PREFIXES_COPY = { "datadog.slf4j", + "datadog.context", "datadog.appsec.api", "datadog.trace.api", "datadog.trace.bootstrap", diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle index 24b06c8351b..2fd95f2c4af 100644 --- a/gradle/dependencies.gradle +++ b/gradle/dependencies.gradle @@ -16,6 +16,7 @@ final class CachedData { exclude(project(':internal-api')) exclude(project(':internal-api:internal-api-9')) exclude(project(':communication')) + exclude(project(':components:context')) exclude(project(':components:json')) exclude(project(':remote-config:remote-config-api')) exclude(project(':remote-config:remote-config-core')) diff --git a/internal-api/build.gradle b/internal-api/build.gradle index e0da08b57fa..fab211ca323 100644 --- a/internal-api/build.gradle +++ b/internal-api/build.gradle @@ -81,6 +81,7 @@ excludedClassesCoverage += [ "datadog.trace.bootstrap.instrumentation.api.StatsPoint", "datadog.trace.bootstrap.instrumentation.api.Schema", "datadog.trace.bootstrap.instrumentation.api.ScopeSource", + "datadog.trace.bootstrap.instrumentation.api.InternalContextKeys", "datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes", "datadog.trace.bootstrap.instrumentation.api.TagContext", "datadog.trace.bootstrap.instrumentation.api.TagContext.HttpHeaders", @@ -210,6 +211,7 @@ dependencies { // references TraceScope and Continuation from public api api project(':dd-trace-api') api libs.slf4j + api project(':components:context') api project(":utils:time-utils") // has to be loaded by system classloader: diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java index a9ca2e11178..0df378b912c 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java @@ -1,12 +1,19 @@ package datadog.trace.bootstrap.instrumentation.api; +import static datadog.trace.bootstrap.instrumentation.api.InternalContextKeys.SPAN_KEY; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.context.ImplicitContextKeyed; import datadog.trace.api.DDTraceId; import datadog.trace.api.TraceConfig; import datadog.trace.api.gateway.IGSpanInfo; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.interceptor.MutableSpan; +import javax.annotation.Nullable; -public interface AgentSpan extends MutableSpan, IGSpanInfo, WithAgentSpan { +public interface AgentSpan + extends MutableSpan, ImplicitContextKeyed, Context, IGSpanInfo, WithAgentSpan { DDTraceId getTraceId(); @@ -145,4 +152,21 @@ public interface AgentSpan extends MutableSpan, IGSpanInfo, WithAgentSpan { default AgentSpan asAgentSpan() { return this; } + + @Override + default Context storeInto(Context context) { + return context.with(SPAN_KEY, this); + } + + @Nullable + @Override + default T get(ContextKey key) { + // noinspection unchecked + return SPAN_KEY == key ? (T) this : Context.root().get(key); + } + + @Override + default Context with(ContextKey key, @Nullable T value) { + return Context.root().with(SPAN_KEY, this, key, value); + } } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InternalContextKeys.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InternalContextKeys.java new file mode 100644 index 00000000000..296b45fec23 --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InternalContextKeys.java @@ -0,0 +1,9 @@ +package datadog.trace.bootstrap.instrumentation.api; + +import datadog.context.ContextKey; + +final class InternalContextKeys { + static final ContextKey SPAN_KEY = ContextKey.named("dd-span-key"); + + private InternalContextKeys() {} +}