Implement Iterator.remove() on OptimizedTagMap view iterators

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

@@ -2082,16 +2082,19 @@ String toInternalString() {
   }
 
   abstract static class IteratorBase {
+    private final OptimizedTagMap map;
     private final Object[] buckets;
 
     private Entry nextEntry;
+    private Entry lastReturnedEntry;
 
     private int bucketIndex = -1;
 
     private BucketGroup group = null;
     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() {

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<TagMap.EntryReader> 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<TagMap.EntryReader> 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<TagMap.EntryReader> 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<TagMap.EntryReader> 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<TagMap.EntryReader> iter = map.iterator();
+    iter.next();
+    map.freeze();
+    assertThrows(IllegalStateException.class, iter::remove);
+  }
+
   @ParameterizedTest
   @EnumSource(TagMapScenario.class)
   public void map_put(TagMapScenario scenario) {