[bidi][java] Ensure empty values are removed from the call back map

NDViet · Jul 3, 2024 · 06e7b7a · 06e7b7a
1 parent 5c08747
commit 06e7b7a
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/java/src/org/openqa/selenium/bidi/Connection.java b/java/src/org/openqa/selenium/bidi/Connection.java
@@ -25,6 +25,7 @@
 import java.io.Closeable;
 import java.io.StringReader;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -210,7 +211,16 @@ public void removeListener(long id) {
     Lock lock = callbacksLock.writeLock();
     lock.lock();
     try {
-      eventCallbacks.forEach((k, v) -> v.remove(id));
+      List<Event<?>> list = new ArrayList<>();
+      eventCallbacks.forEach(
+          (k, v) -> {
+            v.remove(id);
+            if (v.isEmpty()) {
+              list.add(k);
+            }
+          });
+
+      list.forEach(eventCallbacks::remove);
     } finally {
       lock.unlock();
     }