[bidi] [java] Ensure the listeners returns an id (#14215)

Author: pujagani

java/src/org/openqa/selenium/bidi/BiDi.java

@@ -51,15 +51,15 @@ public <X> X send(Command<X> command) {
     return connection.sendAndWait(command, timeout);
   }
 
-  public <X> void addListener(Event<X> event, Consumer<X> handler) {
+  public <X> long addListener(Event<X> event, Consumer<X> handler) {
     Require.nonNull("Event to listen for", event);
     Require.nonNull("Handler to call", handler);
 
     send(
         new Command<>(
             "session.subscribe", Map.of("events", Collections.singletonList(event.getMethod()))));
 
-    connection.addListener(event, handler);
+    return connection.addListener(event, handler);
   }
 
   public <X> void addListener(String browsingContextId, Event<X> event, Consumer<X> handler) {
@@ -79,7 +79,7 @@ public <X> void addListener(String browsingContextId, Event<X> event, Consumer<X
     connection.addListener(event, handler);
   }
 
-  public <X> void addListener(Set<String> browsingContextIds, Event<X> event, Consumer<X> handler) {
+  public <X> long addListener(Set<String> browsingContextIds, Event<X> event, Consumer<X> handler) {
     Require.nonNull("List of browsing context ids", browsingContextIds);
     Require.nonNull("Event to listen for", event);
     Require.nonNull("Handler to call", handler);
@@ -93,7 +93,7 @@ public <X> void addListener(Set<String> browsingContextIds, Event<X> event, Cons
                 "events",
                 Collections.singletonList(event.getMethod()))));
 
-    connection.addListener(event, handler);
+    return connection.addListener(event, handler);
   }
 
   public <X> void clearListener(Event<X> event) {
@@ -111,6 +111,10 @@ public <X> void clearListener(Event<X> event) {
     }
   }
 
+  public void removeListener(long id) {
+    connection.removeListener(id);
+  }
+
   public void clearListeners() {
     connection.clearListeners();
   }

java/src/org/openqa/selenium/bidi/Connection.java

@@ -25,7 +25,6 @@
 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;
@@ -66,11 +65,12 @@ public class Connection implements Closeable {
             return thread;
           });
   private static final AtomicLong NEXT_ID = new AtomicLong(1L);
+  private final AtomicLong EVENT_CALLBACK_ID = new AtomicLong(1);
   private WebSocket socket;
   private final Map<Long, Consumer<Either<Throwable, JsonInput>>> methodCallbacks =
       new ConcurrentHashMap<>();
   private final ReadWriteLock callbacksLock = new ReentrantReadWriteLock(true);
-  private final Map<Event<?>, List<Consumer<?>>> eventCallbacks = new HashMap<>();
+  private final Map<Event<?>, Map<Long, Consumer<?>>> eventCallbacks = new HashMap<>();
   private final HttpClient client;
   private final AtomicBoolean underlyingSocketClosed = new AtomicBoolean();
 
@@ -180,17 +180,26 @@ public <X> X sendAndWait(Command<X> command, Duration timeout) {
     }
   }
 
-  public <X> void addListener(Event<X> event, Consumer<X> handler) {
+  public <X> long addListener(Event<X> event, Consumer<X> handler) {
     Require.nonNull("Event to listen for", event);
     Require.nonNull("Handler to call", handler);
 
+    long id = EVENT_CALLBACK_ID.getAndIncrement();
+
     Lock lock = callbacksLock.writeLock();
     lock.lock();
     try {
-      eventCallbacks.computeIfAbsent(event, (key) -> new ArrayList<>()).add(handler);
+      eventCallbacks.computeIfAbsent(
+          event,
+          key -> {
+            HashMap<Long, Consumer<?>> map = new HashMap<>();
+            map.put(id, handler);
+            return map;
+          });
     } finally {
       lock.unlock();
     }
+    return id;
   }
 
   public <X> void clearListener(Event<X> event) {
@@ -203,6 +212,23 @@ public <X> void clearListener(Event<X> event) {
     }
   }
 
+  public void removeListener(long id) {
+    Lock lock = callbacksLock.writeLock();
+    lock.lock();
+    try {
+      eventCallbacks.forEach((k, v) -> v.remove(id));
+      eventCallbacks.forEach(
+          (k, v) -> {
+            v.remove(id);
+            if (v.isEmpty()) {
+              eventCallbacks.remove(k);
+            }
+          });
+    } finally {
+      lock.unlock();
+    }
+  }
+
   public <X> boolean isEventSubscribed(Event<X> event) {
     Lock lock = callbacksLock.writeLock();
     lock.lock();
@@ -354,7 +380,7 @@ private void handleEventResponse(Map<String, Object> rawDataMap) {
 
                 final Object finalValue = value;
 
-                for (Consumer<?> action : event.getValue()) {
+                for (Consumer<?> action : event.getValue().values()) {
                   @SuppressWarnings("unchecked")
                   Consumer<Object> obj = (Consumer<Object>) action;
                   LOG.log(

java/src/org/openqa/selenium/bidi/module/LogInspector.java

@@ -61,11 +61,11 @@ public LogInspector(Set<String> browsingContextIds, WebDriver driver) {
     this.browsingContextIds = browsingContextIds;
   }
 
-  public void onConsoleEntry(Consumer<ConsoleLogEntry> consumer) {
+  public long onConsoleEntry(Consumer<ConsoleLogEntry> consumer) {
     Consumer<LogEntry> logEntryConsumer =
         logEntry -> logEntry.getConsoleLogEntry().ifPresent(consumer);
 
-    addLogEntryAddedListener(logEntryConsumer);
+    return addLogEntryAddedListener(logEntryConsumer);
   }
 
   public void onConsoleEntry(Consumer<ConsoleLogEntry> consumer, FilterBy filter) {
@@ -105,7 +105,7 @@ public void onJavaScriptLog(Consumer<JavascriptLogEntry> consumer, FilterBy filt
     addLogEntryAddedListener(logEntryConsumer);
   }
 
-  public void onJavaScriptException(Consumer<JavascriptLogEntry> consumer) {
+  public long onJavaScriptException(Consumer<JavascriptLogEntry> consumer) {
     Consumer<LogEntry> logEntryConsumer =
         logEntry ->
             logEntry
@@ -117,7 +117,7 @@ public void onJavaScriptException(Consumer<JavascriptLogEntry> consumer) {
                       }
                     });
 
-    addLogEntryAddedListener(logEntryConsumer);
+    return addLogEntryAddedListener(logEntryConsumer);
   }
 
   public void onGenericLog(Consumer<GenericLogEntry> consumer) {
@@ -163,11 +163,11 @@ public void onLog(Consumer<LogEntry> consumer, FilterBy filter) {
     addLogEntryAddedListener(logEntryConsumer);
   }
 
-  private void addLogEntryAddedListener(Consumer<LogEntry> consumer) {
+  private long addLogEntryAddedListener(Consumer<LogEntry> consumer) {
     if (browsingContextIds.isEmpty()) {
-      this.bidi.addListener(Log.entryAdded(), consumer);
+      return this.bidi.addListener(Log.entryAdded(), consumer);
     } else {
-      this.bidi.addListener(browsingContextIds, Log.entryAdded(), consumer);
+      return this.bidi.addListener(browsingContextIds, Log.entryAdded(), consumer);
     }
   }