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

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java

@@ -167,6 +167,6 @@ default <T> T get(ContextKey<T> key) {
 
   @Override
   default <T> Context with(ContextKey<T> key, @Nullable T value) {
-    return Context.root().with(SPAN_KEY, this).with(key, value);
+    return Context.root().with(SPAN_KEY, this, key, value);
   }
 }