Allow use of span as a primordial context (#8250)

* Support Context.with(span) for convenience
* Allow AgentSpan to act as a primordial Context
* Allow context implementation to skip interim copy when adding two entries at the same time (such as adding span plus baggage)

Author: mcculls

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.
    *
-   * <p>Existing value with the given key will be replaced, and mapping to a {@code null} value will
+   * <p>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 <T> the type of the value.
@@ -124,6 +124,28 @@ static Context detachFrom(Object carrier) {
    */
   <T> Context with(ContextKey<T> key, @Nullable T value);
 
+  /**
+   * Creates a copy of this context with the given pair of key-values.
+   *
+   * <p>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 <T> the type of the first value.
+   * @param <U> 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 <T, U> Context with(
+      ContextKey<T> firstKey,
+      @Nullable T firstValue,
+      ContextKey<U> 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.
    *

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) {

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",

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::'))
   }

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",

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'))

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:

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> T get(ContextKey<T> key) {
+    // noinspection unchecked
+    return SPAN_KEY == key ? (T) this : Context.root().get(key);
+  }
+
+  @Override
+  default <T> Context with(ContextKey<T> key, @Nullable T value) {
+    return Context.root().with(SPAN_KEY, this, key, value);
+  }
 }

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<AgentSpan> SPAN_KEY = ContextKey.named("dd-span-key");
+
+  private InternalContextKeys() {}
+}