Abstract OT-Baggage Tests

dd-trace-core/src/test/java/datadog/trace/core/propagation/AbstractOTBaggageTest.java

@@ -0,0 +1,111 @@
+package datadog.trace.core.propagation;
+
+import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_BYTES;
+import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_ITEMS;
+import static datadog.trace.bootstrap.instrumentation.api.ContextVisitors.stringValuesMap;
+import static datadog.trace.core.propagation.HttpCodecTestHelper.baggageItems;
+import static datadog.trace.core.propagation.HttpCodecTestHelper.generateBaggageItems;
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.singletonMap;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import datadog.trace.bootstrap.instrumentation.api.TagContext;
+import datadog.trace.test.junit.utils.config.WithConfig;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Behavior contract for the baggage limits enforced by {@link ContextInterpreter#addBaggageItem}.
+ * The limits are applied after the propagation style has parsed the key and value, so they are
+ * style independent: every codec carrying caller baggage runs this suite as a {@code @Nested} class
+ * of its extractor test, supplying only the wire format.
+ */
+abstract class AbstractOTBaggageTest {
+
+  /** Returns the extractor under test, built by the enclosing extractor test. */
+  protected abstract HttpCodec.Extractor extractor();
+
+  /**
+   * Returns the given baggage items, in order, in this style's wire format. A repeated key must be
+   * spelled so that the style reads it back as the same baggage key.
+   */
+  protected abstract Map<String, String> baggageHeaders(List<Entry<String, String>> items);
+
+  /** Returns the headers carrying {@code itemCount} generated baggage items. */
+  protected final Map<String, String> generateBaggageHeaders(int itemCount) {
+    return baggageHeaders(generateBaggageItems(itemCount));
+  }
+
+  /**
+   * Extracts the given headers and returns the baggage of the resulting context. Headers carrying
+   * nothing but baggage create no context at all once every item is dropped, which is reported as
+   * empty baggage rather than as a missing context.
+   */
+  protected final Map<String, String> extractBaggage(Map<String, String> headers) {
+    TagContext context = this.extractor().extract(headers, stringValuesMap());
+    return context == null ? emptyMap() : context.getBaggage();
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "3")
+  void stopsAtItemLimit() {
+    assertEquals(3, extractBaggage(generateBaggageHeaders(50)).size());
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "24")
+  void stopsAtByteLimit() {
+    // with single digit indices each stored item is "keyN" + "valueN" = 10 bytes, so 2 fit in 24
+    // bytes and a third would take the total to 30
+    assertEquals(2, extractBaggage(generateBaggageHeaders(10)).size());
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "24")
+  @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "2")
+  void chargesRepeatedKeyOnce() {
+    Map<String, String> baggage =
+        extractBaggage(
+            baggageHeaders(baggageItems("key0", "val0", "a", "0123456789", "key0", "val1")));
+
+    Map<String, String> expected = new HashMap<>();
+    expected.put("key0", "val1");
+    expected.put("a", "0123456789");
+    assertEquals(expected, baggage);
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "8")
+  void chargesEncodedValueSize() {
+    // the budget is charged before decoding, so the 9 character raw value does not fit
+    Map<String, String> baggage =
+        extractBaggage(baggageHeaders(baggageItems("a", "b", "c", "%E2%99%A5")));
+
+    assertEquals(singletonMap("a", "b"), baggage);
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "3")
+  void chargesLiteralUtf8ByCharacterCount() {
+    // "a" + "♥" is 2 characters, so the 2 characters of "b" + "c" no longer fit
+    Map<String, String> baggage = extractBaggage(baggageHeaders(baggageItems("a", "♥", "b", "c")));
+
+    assertEquals(singletonMap("a", "♥"), baggage);
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "0")
+  void dropsAllBaggageWhenItemLimitIsZero() {
+    assertTrue(extractBaggage(generateBaggageHeaders(10)).isEmpty());
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "0")
+  void dropsAllBaggageWhenByteLimitIsZero() {
+    assertTrue(extractBaggage(generateBaggageHeaders(10)).isEmpty());
+  }
+}

dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogHttpExtractorTest.java

@@ -1,7 +1,6 @@
 package datadog.trace.core.propagation;
 
 import static datadog.trace.api.config.TracerConfig.REQUEST_HEADER_TAGS_COMMA_ALLOWED;
-import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_BYTES;
 import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_ITEMS;
 import static datadog.trace.api.sampling.PrioritySampling.UNSET;
 import static datadog.trace.bootstrap.instrumentation.api.ContextVisitors.stringValuesMap;
@@ -12,6 +11,7 @@
 import static datadog.trace.core.propagation.DatadogHttpCodec.SPAN_ID_KEY;
 import static datadog.trace.core.propagation.DatadogHttpCodec.TRACE_ID_KEY;
 import static datadog.trace.core.propagation.HttpCodecTestHelper.headers;
+import static datadog.trace.core.propagation.HttpCodecTestHelper.otBaggageHeaders;
 import static datadog.trace.test.junit.utils.converter.TraceIdConverter.TRACE_ID_MAX_PLUS_1;
 import static java.util.Collections.singletonMap;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -34,8 +34,11 @@
 import datadog.trace.test.junit.utils.converter.TraceIdConverter;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.function.Supplier;
+import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.converter.ConvertWith;
@@ -328,117 +331,31 @@ void baggageIsMappedOnContextCreation(
     }
   }
 
-  @Test
-  @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "3")
-  void extractOtBaggageStopsAtItemLimit() {
-    Map<String, String> headers = otBaggageHeaders(50);
-    headers.put(SOME_CUSTOM_BAGGAGE_HEADER, "mappedBaggageValue");
-
-    TagContext context = this.extractor.extract(headers, stringValuesMap());
-
-    assertEquals(3, context.getBaggage().size());
-  }
-
-  @Test
-  @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "1")
-  void extractMappedBaggageStopsAtItemLimit() {
-    Map<String, String> headers = otBaggageHeaders(50);
-    headers.put(SOME_CUSTOM_BAGGAGE_HEADER, "mappedBaggageValue");
-
-    TagContext context = this.extractor.extract(headers, stringValuesMap());
-
-    assertEquals(1, context.getBaggage().size());
-  }
-
-  @Test
-  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "24")
-  void extractOtBaggageStopsAtByteLimit() {
-    // with single digit indices each stored item is "keyN" + "valueN" = 10 bytes, so 2 fit in 24
-    // bytes and a third would take the total to 30
-    TagContext context = this.extractor.extract(otBaggageHeaders(10), stringValuesMap());
-
-    assertEquals(2, context.getBaggage().size());
-  }
-
-  @Test
-  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "24")
-  void extractOtBaggageChargesRepeatedKeyOnce() {
-    // headers are visited in insertion order, so the duplicate key is seen before the last item
-    Map<String, String> headers = new LinkedHashMap<>();
-    // "key0" + "val0" is 8 bytes, and the duplicate refunds the value it replaces, so the total
-    // stays at 8 rather than doubling
-    headers.put(OT_BAGGAGE_PREFIX + "key0", "val0");
-    headers.put(OT_BAGGAGE_PREFIX + "KEY0", "val0");
-    // leaving room for these 11 bytes, taking the total to 19
-    headers.put(OT_BAGGAGE_PREFIX + "a", "0123456789");
-
-    TagContext context = this.extractor.extract(headers, stringValuesMap());
-
-    Map<String, String> expected = new HashMap<>();
-    expected.put("key0", "val0");
-    expected.put("a", "0123456789");
-    assertEquals(expected, context.getBaggage());
-  }
-
   @Test
   @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "1")
-  void extractOtBaggageAllowsReplacementAtItemLimit() {
-    Map<String, String> headers = new LinkedHashMap<>();
-    headers.put(OT_BAGGAGE_PREFIX + "key0", "old");
-    headers.put(OT_BAGGAGE_PREFIX + "KEY0", "replacement");
-
-    TagContext context = this.extractor.extract(headers, stringValuesMap());
-
-    assertEquals(singletonMap("key0", "replacement"), context.getBaggage());
-  }
-
-  @Test
-  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "24")
-  void extractOtBaggageChargesOnlyTheDeltaWhenReplacingAValue() {
-    Map<String, String> headers = new LinkedHashMap<>();
-    headers.put(OT_BAGGAGE_PREFIX + "key0", "val0"); // 8 bytes
-    // replaces the value, charging the 8 byte difference rather than another 16 bytes
-    headers.put(OT_BAGGAGE_PREFIX + "KEY0", "012345678901");
-    headers.put(OT_BAGGAGE_PREFIX + "a", "0123456"); // 8 bytes, taking the total to exactly 24
-
-    TagContext context = this.extractor.extract(headers, stringValuesMap());
-
-    Map<String, String> expected = new HashMap<>();
-    expected.put("key0", "012345678901");
-    expected.put("a", "0123456");
-    assertEquals(expected, context.getBaggage());
-  }
-
-  @Test
-  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "8")
-  void extractOtBaggageChargesEncodedValueSize() {
+  void extractMappedBaggageIsSubjectToTheItemLimit() {
+    // mapped baggage shares the item budget with the baggage read off the wire, so the wire item
+    // is dropped once the mapped header has claimed the only slot
     Map<String, String> headers = new LinkedHashMap<>();
-    headers.put(OT_BAGGAGE_PREFIX + "a", "b"); // 2 characters
-    headers.put(OT_BAGGAGE_PREFIX + "c", "%E2%99%A5"); // 1 character key + 9 character raw value
+    headers.put(SOME_CUSTOM_BAGGAGE_HEADER, "mappedBaggageValue");
+    headers.put(OT_BAGGAGE_PREFIX + "wireKey", "wireValue");
 
     TagContext context = this.extractor.extract(headers, stringValuesMap());
 
-    assertEquals(singletonMap("a", "b"), context.getBaggage());
+    assertEquals(singletonMap(SOME_BAGGAGE, "mappedBaggageValue"), context.getBaggage());
   }
 
-  @Test
-  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "3")
-  void extractOtBaggageChargesLiteralUtf8ByCharacterCount() {
-    Map<String, String> headers = new LinkedHashMap<>();
-    headers.put(OT_BAGGAGE_PREFIX + "a", "♥"); // 2 characters
-    headers.put(OT_BAGGAGE_PREFIX + "b", "c"); // 2 more characters, no longer fits
-
-    TagContext context = this.extractor.extract(headers, stringValuesMap());
-
-    assertEquals(singletonMap("a", "♥"), context.getBaggage());
-  }
+  @Nested
+  class BaggageLimits extends AbstractOTBaggageTest {
+    @Override
+    protected HttpCodec.Extractor extractor() {
+      return DatadogHttpExtractorTest.this.extractor;
+    }
 
-  private static Map<String, String> otBaggageHeaders(int count) {
-    Map<String, String> headers = new HashMap<>();
-    for (int i = 0; i < count; i++) {
-      headers.put(OT_BAGGAGE_PREFIX + "key" + i, "value" + i);
+    @Override
+    protected Map<String, String> baggageHeaders(List<Entry<String, String>> items) {
+      return otBaggageHeaders(OT_BAGGAGE_PREFIX, items);
     }
-    return headers;
   }
 
   private static String asString(CharSequence cs) {

dd-trace-core/src/test/java/datadog/trace/core/propagation/HaystackHttpExtractorTest.java

@@ -10,7 +10,9 @@
 import static datadog.trace.core.propagation.HaystackHttpCodec.PARENT_ID_KEY;
 import static datadog.trace.core.propagation.HaystackHttpCodec.SPAN_ID_KEY;
 import static datadog.trace.core.propagation.HaystackHttpCodec.TRACE_ID_KEY;
+import static datadog.trace.core.propagation.HttpCodecTestHelper.generateBaggageItems;
 import static datadog.trace.core.propagation.HttpCodecTestHelper.headers;
+import static datadog.trace.core.propagation.HttpCodecTestHelper.otBaggageHeaders;
 import static datadog.trace.test.junit.utils.converter.TraceIdConverter.TRACE_ID_MAX_PLUS_1;
 import static java.util.Collections.singletonMap;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -26,8 +28,11 @@
 import datadog.trace.test.junit.utils.config.WithConfig;
 import datadog.trace.test.junit.utils.converter.TraceIdConverter;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.function.Supplier;
+import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.converter.ConvertWith;
 import org.tabletest.junit.TableTest;
@@ -39,28 +44,26 @@ protected HttpCodec.Extractor newExtractor(
     return HaystackHttpCodec.newExtractor(config, traceConfigSupplier);
   }
 
-  @Test
-  @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "3")
-  void extractBaggageStopsAtItemLimit() {
-    Map<String, String> headers = new HashMap<>();
-    for (int i = 0; i < 50; i++) {
-      headers.put(OT_BAGGAGE_PREFIX + "key" + i, "value" + i);
+  @Nested
+  class BaggageLimits extends AbstractOTBaggageTest {
+    @Override
+    protected HttpCodec.Extractor extractor() {
+      return HaystackHttpExtractorTest.this.extractor;
     }
 
-    TagContext context = this.extractor.extract(headers, stringValuesMap());
-
-    assertEquals(3, context.getBaggage().size());
+    @Override
+    protected Map<String, String> baggageHeaders(List<Entry<String, String>> items) {
+      return otBaggageHeaders(OT_BAGGAGE_PREFIX, items);
+    }
   }
 
   @Test
   @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "1")
   void extractKeepsHaystackIdsWhenBaggageLimitReached() {
     // the Haystack ids are recorded by the tracer for lossless injection, so caller supplied
     // Baggage-* headers must not be able to evict them by exhausting the item limit
-    Map<String, String> headers = new HashMap<>();
-    for (int i = 0; i < 50; i++) {
-      headers.put(OT_BAGGAGE_PREFIX + "key" + i, "value" + i);
-    }
+    Map<String, String> headers =
+        new HashMap<>(otBaggageHeaders(OT_BAGGAGE_PREFIX, generateBaggageItems(50)));
     headers.put(TRACE_ID_KEY, "44617461-646f-6721-0000-000000000001");
     headers.put(SPAN_ID_KEY, "44617461-646f-6721-0000-000000000002");