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() {}
+}