From f739c89af4ae5e73e0f13ac563a414fab1b02440 Mon Sep 17 00:00:00 2001 From: Douglas Q Hawkins Date: Fri, 8 May 2026 10:51:15 -0400 Subject: [PATCH] Implement Iterator.remove() on OptimizedTagMap view iterators The view iterators returned by OptimizedTagMap.entrySet().iterator(), keySet().iterator(), values().iterator(), and the TagMap.iterator() overload all extend IteratorBase, which inherited the default Iterator.remove() that throws UnsupportedOperationException. As a result, AbstractSet/AbstractCollection bulk operations that rely on iterator-remove (removeAll, removeIf, retainAll) failed mid-iteration on optimized TagMaps. Track the last-returned entry on IteratorBase and remove by tag through OptimizedTagMap.getAndRemove, which enforces the frozen check. LegacyTagMap is unaffected; its view iterators are inherited from HashMap and already support remove. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../main/java/datadog/trace/api/TagMap.java | 29 ++++- .../java/datadog/trace/api/TagMapTest.java | 120 ++++++++++++++++++ 2 files changed, 144 insertions(+), 5 deletions(-) diff --git a/internal-api/src/main/java/datadog/trace/api/TagMap.java b/internal-api/src/main/java/datadog/trace/api/TagMap.java index 95f676245ef..49223c5b53b 100644 --- a/internal-api/src/main/java/datadog/trace/api/TagMap.java +++ b/internal-api/src/main/java/datadog/trace/api/TagMap.java @@ -2082,9 +2082,11 @@ String toInternalString() { } abstract static class IteratorBase { + private final OptimizedTagMap map; private final Object[] buckets; private Entry nextEntry; + private Entry lastReturnedEntry; private int bucketIndex = -1; @@ -2092,6 +2094,7 @@ abstract static class IteratorBase { private int groupIndex = 0; IteratorBase(OptimizedTagMap map) { + this.map = map; this.buckets = map.buckets; } @@ -2108,12 +2111,14 @@ public final boolean hasNext() { final Entry nextEntryOrThrowNoSuchElement() { if (this.nextEntry != null) { - Entry nextEntry = this.nextEntry; + Entry result = this.nextEntry; this.nextEntry = null; - return nextEntry; + this.lastReturnedEntry = result; + return result; } if (this.hasNext()) { + this.lastReturnedEntry = this.nextEntry; return this.nextEntry; } else { throw new NoSuchElementException(); @@ -2122,12 +2127,26 @@ final Entry nextEntryOrThrowNoSuchElement() { final Entry nextEntryOrNull() { if (this.nextEntry != null) { - Entry nextEntry = this.nextEntry; + Entry result = this.nextEntry; this.nextEntry = null; - return nextEntry; + this.lastReturnedEntry = result; + return result; } - return this.hasNext() ? this.nextEntry : null; + if (this.hasNext()) { + this.lastReturnedEntry = this.nextEntry; + return this.nextEntry; + } + return null; + } + + public final void remove() { + Entry last = this.lastReturnedEntry; + if (last == null) { + throw new IllegalStateException("next() must be called before remove()"); + } + this.lastReturnedEntry = null; + this.map.getAndRemove(last.tag); } private final Entry advance() { diff --git a/internal-api/src/test/java/datadog/trace/api/TagMapTest.java b/internal-api/src/test/java/datadog/trace/api/TagMapTest.java index 365ccff23d8..6fb2c9f6020 100644 --- a/internal-api/src/test/java/datadog/trace/api/TagMapTest.java +++ b/internal-api/src/test/java/datadog/trace/api/TagMapTest.java @@ -5,8 +5,10 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.HashMap; @@ -94,6 +96,124 @@ public void optimizedFactory(boolean optimized) { assertEquals(optimized, emptyMap.isOptimized()); } + @ParameterizedTest + @EnumSource(TagMapType.class) + public void entrySet_removeIf(TagMapType type) { + TagMap map = type.create(); + map.set("a", 1); + map.set("b", 2); + map.set("c", 3); + + boolean removed = + map.entrySet().removeIf(e -> "a".equals(e.getKey()) || "c".equals(e.getKey())); + + assertTrue(removed); + assertEquals(1, map.size()); + assertNull(map.getEntry("a")); + assertNotNull(map.getEntry("b")); + assertNull(map.getEntry("c")); + } + + @ParameterizedTest + @EnumSource(TagMapType.class) + public void keySet_removeAll(TagMapType type) { + TagMap map = type.create(); + map.set("a", 1); + map.set("b", 2); + map.set("c", 3); + + boolean removed = map.keySet().removeAll(Arrays.asList("a", "c")); + + assertTrue(removed); + assertEquals(1, map.size()); + assertNull(map.getEntry("a")); + assertNotNull(map.getEntry("b")); + assertNull(map.getEntry("c")); + } + + @ParameterizedTest + @EnumSource(TagMapType.class) + public void values_removeIf(TagMapType type) { + TagMap map = type.create(); + map.set("a", 1); + map.set("b", 2); + map.set("c", 3); + + boolean removed = map.values().removeIf(v -> v.equals(1) || v.equals(3)); + + assertTrue(removed); + assertEquals(1, map.size()); + assertNull(map.getEntry("a")); + assertNotNull(map.getEntry("b")); + assertNull(map.getEntry("c")); + } + + @Test + public void optimizedIterator_remove() { + TagMap map = TagMapType.OPTIMIZED.create(); + map.set("a", 1); + map.set("b", 2); + + Iterator iter = map.iterator(); + while (iter.hasNext()) { + if ("a".equals(iter.next().tag())) { + iter.remove(); + } + } + + assertEquals(1, map.size()); + assertNull(map.getEntry("a")); + assertNotNull(map.getEntry("b")); + } + + @Test + public void optimizedIterator_remove_acrossBucketCollisions() { + // OptimizedTagMap uses 16 buckets, so 64 entries guarantees BucketGroup chains. + TagMap map = TagMapType.OPTIMIZED.create(); + for (int i = 0; i < 64; i++) { + map.set("k" + i, i); + } + assertEquals(64, map.size()); + + Iterator iter = map.iterator(); + while (iter.hasNext()) { + iter.next(); + iter.remove(); + } + + assertEquals(0, map.size()); + assertTrue(map.isEmpty()); + } + + @Test + public void optimizedIterator_remove_beforeNext_throws() { + TagMap map = TagMapType.OPTIMIZED.create(); + map.set("a", 1); + Iterator iter = map.iterator(); + assertThrows(IllegalStateException.class, iter::remove); + } + + @Test + public void optimizedIterator_remove_twice_throws() { + TagMap map = TagMapType.OPTIMIZED.create(); + map.set("a", 1); + Iterator iter = map.iterator(); + iter.next(); + iter.remove(); + assertThrows(IllegalStateException.class, iter::remove); + } + + @Test + public void optimizedIterator_remove_onFrozenMap_throws() { + TagMap map = TagMapType.OPTIMIZED.create(); + map.set("a", 1); + map.set("b", 2); + Iterator iter = map.iterator(); + iter.next(); + map.freeze(); + assertThrows(IllegalStateException.class, iter::remove); + } + @ParameterizedTest @EnumSource(TagMapScenario.class) public void map_put(TagMapScenario scenario) {