Make a null tag value remove the tag in TagMap (dense store - part 1)

dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/PayloadTagsProcessor.java

@@ -245,6 +245,7 @@ public void doubleValue(PathCursor path, double value) {
 
     @Override
     public void nullValue(PathCursor path) {
+      // a JSON null sets no tag, and removes an earlier value set at the same path
       collectedTags.set(path.toString(tagPrefix), null);
     }
 

dd-trace-core/src/test/java/datadog/trace/core/tagprocessor/PayloadTagsProcessorTest.java

@@ -203,8 +203,6 @@ void expandPreservingTagTypes() {
             3.14d,
             "payload.tag4.Value.0",
             "string",
-            "payload.tag5",
-            null,
             "payload.tag6",
             false);
     assertEquals(expected, unsafeTags);
@@ -246,8 +244,6 @@ void expandStringifiedJsonTags() {
             2,
             "p.j3.2",
             3.14d,
-            "p.j3.3",
-            null,
             "p.j3.4",
             true,
             "p.j4.foo",
@@ -257,6 +253,23 @@ void expandStringifiedJsonTags() {
     assertEquals(expected, unsafeTags);
   }
 
+  @Test
+  void nullRemovesEarlierValueAtSamePath() {
+    PayloadTagsProcessor ptp = tagsProcessor("p", Collections.emptyList(), 10, 758);
+
+    Map<String, Object> st =
+        spanTags(
+            "p",
+            pv(pc().push("j1"), "{'x': 'old', 'x': null, 'y': 1}"),
+            pv(pc().push("v"), "old"),
+            pv(pc().push("v"), null));
+
+    TagMap unsafeTags = TagMap.fromMap(st);
+    ptp.processTags(unsafeTags, null, link -> {});
+
+    assertEquals(mapOf("p.j1.y", 1), unsafeTags);
+  }
+
   @Test
   void expandSerializedEscapedInnerJsonWithinInnerJson() throws IOException {
     Buffer b0 = new Buffer();
@@ -410,8 +423,6 @@ void applyRedactionRules() {
             2,
             "p.j3.2",
             3.14d,
-            "p.j3.3",
-            null,
             "p.j3.4",
             true,
             "p.j4.foo",
@@ -444,8 +455,8 @@ void respectMaxTagsLimit() {
             2,
             "p.j3.2",
             3.14d,
-            "p.j3.3",
-            null,
+            "p.j3.4",
+            true,
             "_dd.payload_tags_incomplete",
             true),
         unsafeTags);
@@ -474,8 +485,6 @@ void respectMaxDepthLimit() {
             2,
             "p.j3.2",
             3.14d,
-            "p.j3.3",
-            null,
             "p.j3.4",
             true,
             "p.j3.5.0",

internal-api/src/main/java/datadog/trace/api/TagMap.java

@@ -1502,7 +1502,9 @@ boolean parentEntryVisible(Entry parentEntry, TagMap fromAncestor) {
   @Deprecated
   @Override
   public Object put(@Nonnull String tag, Object value) {
-    TagMap.Entry entry = this.getAndSet(anyEntryFor(tag, value));
+    // a null value removes the tag, rather than mapping it to null as HashMap would
+    TagMap.Entry entry =
+        value == null ? this.getAndRemove(tag) : this.getAndSet(anyEntryFor(tag, value));
     return entry == null ? null : entry.objectValue();
   }
 
@@ -1556,7 +1558,12 @@ static Entry doubleEntryFor(String tag, double value) {
     return tagId != 0 ? Entry.newDoubleEntry(tagId, value) : Entry.newDoubleEntry(tag, value);
   }
 
-  public void set(@Nonnull String tag, @Nonnull Object value) {
+  /** A null value removes the tag, as {@code setTag(tag, null)} does on a span. */
+  public void set(@Nonnull String tag, @Nullable Object value) {
+    if (value == null) {
+      this.getAndRemove(tag);
+      return;
+    }
     long tagId = KnownTagCodec.keyOf(tag);
     if (tagId != 0) {
       this.set(tagId, value);
@@ -1565,7 +1572,12 @@ public void set(@Nonnull String tag, @Nonnull Object value) {
     }
   }
 
-  public void set(@Nonnull String tag, @Nonnull CharSequence value) {
+  /** A null value removes the tag, as {@code setTag(tag, null)} does on a span. */
+  public void set(@Nonnull String tag, @Nullable CharSequence value) {
+    if (value == null) {
+      this.getAndRemove(tag);
+      return;
+    }
     long tagId = KnownTagCodec.keyOf(tag);
     if (tagId != 0) {
       this.set(tagId, value);
@@ -1680,11 +1692,21 @@ public void setKnown(long tagId, double value) {
     }
   }
 
-  public void set(long tagId, @Nonnull Object value) {
+  /** A null value removes the tag, as {@code setTag(tag, null)} does on a span. */
+  public void set(long tagId, @Nullable Object value) {
+    if (value == null) {
+      this.getAndRemove(tagId);
+      return;
+    }
     this.putEntry(Entry.newAnyEntry(tagId, value));
   }
 
-  public void set(long tagId, @Nonnull CharSequence value) {
+  /** A null value removes the tag, as {@code setTag(tag, null)} does on a span. */
+  public void set(long tagId, @Nullable CharSequence value) {
+    if (value == null) {
+      this.getAndRemove(tagId);
+      return;
+    }
     this.putEntry(Entry.newObjectEntry(tagId, value));
   }
 
@@ -1797,11 +1819,11 @@ Entry putEntry(@Nonnull Entry newEntry) {
   }
 
   public Entry getAndSet(@Nonnull String tag, Object value) {
-    return this.getAndSet(anyEntryFor(tag, value));
+    return value == null ? this.getAndRemove(tag) : this.getAndSet(anyEntryFor(tag, value));
   }
 
   public Entry getAndSet(@Nonnull String tag, CharSequence value) {
-    return this.getAndSet(objectEntryFor(tag, value));
+    return value == null ? this.getAndRemove(tag) : this.getAndSet(objectEntryFor(tag, value));
   }
 
   public TagMap.Entry getAndSet(@Nonnull String tag, boolean value) {

internal-api/src/test/java/datadog/trace/api/TagMapNullToleranceTest.java

@@ -6,12 +6,16 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import java.util.HashMap;
+import java.util.Map;
 import org.junit.jupiter.api.Test;
 
 /**
  * Pins the Entry-pathway null-tolerance contract: {@code create(...)} returns null for a null or
  * empty value, and the {@code Entry} sinks ({@code set(Entry)} / {@code getAndSet(Entry)}) treat
- * null as a no-op -- so a null/empty value flows through as "no tag" without any caller guarding.
+ * null as a no-op -- so a null/empty value flows through as "no tag" without any caller guarding. A
+ * null value given to a setter removes the tag, as {@code setTag(tag, null)} does on a span.
  */
 class TagMapNullToleranceTest {
 
@@ -63,4 +67,55 @@ void nullOrEmptyFlowsThroughTheEntryPathwayAsNoTag() {
     assertEquals(1, map.size());
     assertTrue(map.containsKey("present"));
   }
+
+  @Test
+  void setNullRemovesAKnownOrCustomTag() {
+    TagMap map = TagMap.create();
+    map.set(Tags.COMPONENT, "okhttp");
+    map.set("custom.tag", "v");
+
+    map.set(Tags.COMPONENT, (Object) null);
+    map.set("custom.tag", (CharSequence) null);
+    assertEquals(0, map.size());
+    assertFalse(map.containsKey(Tags.COMPONENT));
+    assertFalse(map.containsKey("custom.tag"));
+  }
+
+  @Test
+  void setNullByIdRemovesTheTag() {
+    TagMap map = TagMap.create();
+    map.set(KnownTags.COMPONENT_ID, "okhttp");
+    map.set(KnownTags.COMPONENT_ID, (Object) null);
+    assertNull(map.getObject(Tags.COMPONENT));
+
+    map.set(KnownTags.COMPONENT_ID, "okhttp");
+    map.set(KnownTags.COMPONENT_ID, (CharSequence) null);
+    assertEquals(0, map.size());
+  }
+
+  @Test
+  void putAndGetAndSetNullRemoveTheTagAndReturnItsPriorValue() {
+    TagMap map = TagMap.create();
+    map.set("custom.tag", "v");
+    assertEquals("v", map.put("custom.tag", null));
+    assertFalse(map.containsKey("custom.tag"));
+
+    map.set(Tags.COMPONENT, "okhttp");
+    assertEquals("okhttp", map.getAndSet(Tags.COMPONENT, (Object) null).objectValue());
+    assertNull(map.getAndSet(Tags.COMPONENT, (CharSequence) null), "already removed");
+    assertEquals(0, map.size());
+  }
+
+  @Test
+  void putAllWithANullValueLeavesNoTag() {
+    Map<String, Object> source = new HashMap<>();
+    source.put("custom.tag", null);
+    source.put(Tags.COMPONENT, null);
+    source.put("present", "v");
+
+    TagMap map = TagMap.fromMap(source);
+    assertEquals(1, map.size());
+    assertFalse(map.containsKey("custom.tag"));
+    assertFalse(map.containsKey(Tags.COMPONENT));
+  }
 }