[java][bidi] Track events subscriptions using subscription ids

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

@@ -24,8 +24,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.function.Consumer;
 import java.util.logging.Logger;
 import org.openqa.selenium.Beta;
@@ -38,7 +36,6 @@ public class BiDi implements Closeable {
 
   private final Duration timeout;
   private final Connection connection;
-  private final Map<Event<?>, List<Long>> contextListenerIds = new ConcurrentHashMap<>();
 
   /**
    * @deprecated Use constructor with timeout parameter: {@link #BiDi(Connection, Duration)}
@@ -86,84 +83,81 @@ public <X> X send(Command<X> command, Duration timeout) {
     return connection.sendAndWait(command, timeout);
   }
 
-  public <X> long addListener(Event<X> event, Consumer<X> handler) {
+  public <X> String 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", List.of(event.getMethod()))));
-
-    return connection.addListener(event, handler);
+    String subscriptionId = subscribe(Map.of("events", List.of(event.getMethod())));
+    connection.addListener(subscriptionId, event, handler);
+    return subscriptionId;
   }
 
-  public <X> long addListener(String browsingContextId, Event<X> event, Consumer<X> handler) {
+  public <X> String addListener(String browsingContextId, Event<X> event, Consumer<X> handler) {
     Require.nonNull("Event to listen for", event);
     Require.nonNull("Browsing context id", browsingContextId);
     Require.nonNull("Handler to call", handler);
 
-    send(
-        new Command<>(
-            "session.subscribe",
-            Map.of("contexts", List.of(browsingContextId), "events", List.of(event.getMethod()))));
-
-    long id = connection.addListener(event, handler);
-    contextListenerIds.computeIfAbsent(event, k -> new CopyOnWriteArrayList<>()).add(id);
-    return id;
+    String subscriptionId =
+        subscribe(
+            Map.of("contexts", List.of(browsingContextId), "events", List.of(event.getMethod())));
+    connection.addListener(subscriptionId, event, handler);
+    return subscriptionId;
   }
 
-  public <X> long addListener(Set<String> browsingContextIds, Event<X> event, Consumer<X> handler) {
+  public <X> String addListener(
+      Set<String> browsingContextIds, Event<X> event, Consumer<X> handler) {
     Require.nonEmpty("List of browsing context ids", browsingContextIds);
     Require.nonNull("Event to listen for", event);
     Require.nonNull("Handler to call", handler);
 
-    send(
-        new Command<>(
-            "session.subscribe",
-            Map.of("contexts", browsingContextIds, "events", List.of(event.getMethod()))));
-
-    long id = connection.addListener(event, handler);
-    contextListenerIds.computeIfAbsent(event, k -> new CopyOnWriteArrayList<>()).add(id);
-    return id;
+    String subscriptionId =
+        subscribe(Map.of("contexts", browsingContextIds, "events", List.of(event.getMethod())));
+    connection.addListener(subscriptionId, event, handler);
+    return subscriptionId;
   }
 
-  public <X> void clearListener(Event<X> event) {
-    Require.nonNull("Event to listen for", event);
-
-    // The browser throws an error if we try to unsubscribe an event that was not subscribed in the
-    // first place
-    if (connection.isEventSubscribed(event)) {
-      send(new Command<>("session.unsubscribe", Map.of("events", List.of(event.getMethod()))));
-      connection.clearListener(event);
+  // The subscription id returned by the browser is the sole identifier we need to unsubscribe
+  // later: it is unique regardless of whether the subscription was scoped to events, contexts, or
+  // user contexts, so there is no need to separately track how a listener was subscribed.
+  private String subscribe(Map<String, Object> params) {
+    Map<String, Object> result = send(new Command<>("session.subscribe", params, Map.class));
+    Object subscriptionId = result.get("subscription");
+    if (!(subscriptionId instanceof String) || ((String) subscriptionId).isEmpty()) {
+      throw new BiDiException(
+          "session.subscribe did not return a valid subscription id: " + result);
     }
+    return (String) subscriptionId;
   }
 
   public <X> void clearListener(Set<String> browsingContextIds, Event<X> event) {
     Require.nonEmpty("List of browsing context ids", browsingContextIds);
+    // This is not the correct implementation, everything needs to be tired to subscription id.
+    // This is currently used by hand-written BiDi code, soon to be replaced by generated BiDi code.
+    // After that this method will be removed.
+    // This is just to avoid changes to hand-written classes meanwhile.
+    this.clearListener(event);
+  }
+
+  public <X> void clearListener(Event<X> event) {
     Require.nonNull("Event to listen for", event);
 
     // The browser throws an error if we try to unsubscribe an event that was not subscribed in the
     // first place
-    if (!connection.isEventSubscribed(event)) {
-      return;
-    }
-
-    List<Long> ids = contextListenerIds.remove(event);
-    if (ids != null && !ids.isEmpty()) {
-      // Subscription was context-specific: unsubscribe only for those contexts.
-      send(
-          new Command<>(
-              "session.unsubscribe",
-              Map.of("contexts", browsingContextIds, "events", List.of(event.getMethod()))));
-      ids.forEach(connection::removeListener);
-    } else {
-      // Subscription was global: context-specific unsubscription is invalid per the BiDi protocol,
-      // so fall back to a global unsubscription.
+    if (connection.isEventSubscribed(event)) {
       send(new Command<>("session.unsubscribe", Map.of("events", List.of(event.getMethod()))));
       connection.clearListener(event);
     }
   }
 
-  public void removeListener(long id) {
-    connection.removeListener(id);
+  public void removeListener(String id) {
+    Require.nonNull("Listener id", id);
+
+    // The browser throws an error if we try to unsubscribe a subscription id that was not
+    // subscribed in the first place
+    if (connection.isSubscribed(id)) {
+      send(new Command<>("session.unsubscribe", Map.of("subscriptions", List.of(id))));
+      connection.removeListener(id);
+    }
   }
 
   public void clearListeners() {

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

@@ -69,11 +69,10 @@ 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 final Map<Long, Consumer<Either<Throwable, JsonInput>>> methodCallbacks =
       new ConcurrentHashMap<>();
   private final ReadWriteLock callbacksLock = new ReentrantReadWriteLock(true);
-  private final Map<Event<?>, Map<Long, Consumer<?>>> eventCallbacks = new HashMap<>();
+  private final Map<Event<?>, Map<String, Consumer<?>>> eventCallbacks = new HashMap<>();
   private final HttpClient client;
   private final WebSocket socket;
   private final AtomicBoolean underlyingSocketClosed = new AtomicBoolean(false);
@@ -176,20 +175,18 @@ public <X> X sendAndWait(Command<X> command, Duration timeout) {
     }
   }
 
-  public <X> long addListener(Event<X> event, Consumer<X> handler) {
+  public <X> void addListener(String subscriptionId, Event<X> event, Consumer<X> handler) {
+    Require.nonNull("Subscription id", subscriptionId);
     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 HashMap<>()).put(id, handler);
+      eventCallbacks.computeIfAbsent(event, key -> new HashMap<>()).put(subscriptionId, handler);
     } finally {
       lock.unlock();
     }
-    return id;
   }
 
   public <X> void clearListener(Event<X> event) {
@@ -202,7 +199,9 @@ public <X> void clearListener(Event<X> event) {
     }
   }
 
-  public void removeListener(long id) {
+  public void removeListener(String id) {
+    Require.nonNull("Subscription id", id);
+
     Lock lock = callbacksLock.writeLock();
     lock.lock();
     try {
@@ -221,8 +220,18 @@ public void removeListener(long id) {
     }
   }
 
+  public boolean isSubscribed(String id) {
+    Lock lock = callbacksLock.readLock();
+    lock.lock();
+    try {
+      return eventCallbacks.values().stream().anyMatch(listeners -> listeners.containsKey(id));
+    } finally {
+      lock.unlock();
+    }
+  }
+
   public <X> boolean isEventSubscribed(Event<X> event) {
-    Lock lock = callbacksLock.writeLock();
+    Lock lock = callbacksLock.readLock();
     lock.lock();
     try {
       return eventCallbacks.containsKey(event);

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

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

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

@@ -186,7 +186,7 @@ public void onResponseCompleted(Consumer<ResponseDetails> consumer) {
     }
   }
 
-  public long onAuthRequired(Consumer<ResponseDetails> consumer) {
+  public String onAuthRequired(Consumer<ResponseDetails> consumer) {
     if (browsingContextIds.isEmpty()) {
       return this.bidi.addListener(authRequired, consumer);
     } else {

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

@@ -348,7 +348,7 @@ public void removePreloadScript(String id) {
     this.bidi.send(new Command<>("script.removePreloadScript", Map.of("script", id)));
   }
 
-  public long onMessage(Consumer<Message> consumer) {
+  public String onMessage(Consumer<Message> consumer) {
     if (browsingContextIds.isEmpty()) {
       return this.bidi.addListener(messageEvent, consumer);
     } else {

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

@@ -59,15 +59,15 @@ public SpeculationInspector(Set<String> browsingContextIds, WebDriver driver) {
     this.prefetchStatusUpdatedEvent = Speculation.prefetchStatusUpdated();
   }
 
-  public long onPrefetchStatusUpdated(Consumer<PrefetchStatusUpdatedParameters> consumer) {
+  public String onPrefetchStatusUpdated(Consumer<PrefetchStatusUpdatedParameters> consumer) {
     if (browsingContextIds.isEmpty()) {
       return this.bidi.addListener(this.prefetchStatusUpdatedEvent, consumer);
     } else {
       return this.bidi.addListener(browsingContextIds, this.prefetchStatusUpdatedEvent, consumer);
     }
   }
 
-  public void removeListener(long subscriptionId) {
+  public void removeListener(String subscriptionId) {
     this.bidi.removeListener(subscriptionId);
   }
 

java/src/org/openqa/selenium/remote/RemoteScript.java

@@ -62,27 +62,27 @@ public RemoteScript(WebDriver driver) {
   }
 
   @Override
-  public long addConsoleMessageHandler(Consumer<ConsoleLogEntry> consumer) {
+  public String addConsoleMessageHandler(Consumer<ConsoleLogEntry> consumer) {
     return this.logInspector.onConsoleEntry(consumer);
   }
 
   @Override
-  public void removeConsoleMessageHandler(long id) {
+  public void removeConsoleMessageHandler(String id) {
     this.biDi.removeListener(id);
   }
 
   @Override
-  public long addJavaScriptErrorHandler(Consumer<JavascriptLogEntry> consumer) {
+  public String addJavaScriptErrorHandler(Consumer<JavascriptLogEntry> consumer) {
     return this.logInspector.onJavaScriptException(consumer);
   }
 
   @Override
-  public void removeJavaScriptErrorHandler(long id) {
+  public void removeJavaScriptErrorHandler(String id) {
     this.biDi.removeListener(id);
   }
 
   @Override
-  public long addDomMutationHandler(Consumer<DomMutation> consumer) {
+  public String addDomMutationHandler(Consumer<DomMutation> consumer) {
     String scriptValue =
         Read.resourceAsString(getClass(), "/org/openqa/selenium/remote/bidi-mutation-listener.js");
 
@@ -116,7 +116,7 @@ public long addDomMutationHandler(Consumer<DomMutation> consumer) {
   }
 
   @Override
-  public void removeDomMutationHandler(long id) {
+  public void removeDomMutationHandler(String id) {
     this.biDi.removeListener(id);
   }
 

java/src/org/openqa/selenium/remote/Script.java

@@ -26,17 +26,17 @@
 @Beta
 public interface Script {
 
-  long addConsoleMessageHandler(Consumer<ConsoleLogEntry> consumer);
+  String addConsoleMessageHandler(Consumer<ConsoleLogEntry> consumer);
 
-  void removeConsoleMessageHandler(long id);
+  void removeConsoleMessageHandler(String id);
 
-  long addJavaScriptErrorHandler(Consumer<JavascriptLogEntry> consumer);
+  String addJavaScriptErrorHandler(Consumer<JavascriptLogEntry> consumer);
 
-  void removeJavaScriptErrorHandler(long id);
+  void removeJavaScriptErrorHandler(String id);
 
-  long addDomMutationHandler(Consumer<DomMutation> event);
+  String addDomMutationHandler(Consumer<DomMutation> event);
 
-  void removeDomMutationHandler(long id);
+  void removeDomMutationHandler(String id);
 
   String pin(String script);