From 978a2d6afe614c697a840a1e7d61951e2c07187b Mon Sep 17 00:00:00 2001
From: joerg1985 <16140691+joerg1985@users.noreply.github.com>
Date: Fri, 10 May 2024 11:33:50 +0200
Subject: [PATCH] [java] allow a DevTools listener to determinate the order of
 handler calls (#13921)

Co-authored-by: Diego Molina <diemol@users.noreply.github.com>
---
 .../openqa/selenium/devtools/Connection.java  | 17 +++++++------
 .../openqa/selenium/devtools/DevTools.java    | 24 +++++++++++++++++++
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/java/src/org/openqa/selenium/devtools/Connection.java b/java/src/org/openqa/selenium/devtools/Connection.java
index 27f0105cf7dba..1ae1b5b48d595 100644
--- a/java/src/org/openqa/selenium/devtools/Connection.java
+++ b/java/src/org/openqa/selenium/devtools/Connection.java
@@ -43,6 +43,7 @@
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -70,11 +71,12 @@ public class Connection implements Closeable {
             return thread;
           });
   private static final AtomicLong NEXT_ID = new AtomicLong(1L);
+  private static final AtomicLong NEXT_SEQUENCE = new AtomicLong(1L);
   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<?>, List<BiConsumer<Long, ?>>> eventCallbacks = new HashMap<>();
   private HttpClient client;
   private final String url;
   private final AtomicBoolean isClosed;
@@ -196,7 +198,7 @@ public <X> X sendAndWait(SessionID sessionId, Command<X> command, Duration timeo
     }
   }
 
-  public <X> void addListener(Event<X> event, Consumer<X> handler) {
+  public <X> void addListener(Event<X> event, BiConsumer<Long, X> handler) {
     Require.nonNull("Event to listen for", event);
     Require.nonNull("Handler to call", handler);
 
@@ -230,10 +232,11 @@ private class Listener implements WebSocket.Listener {
 
     @Override
     public void onText(CharSequence data) {
+      long sequence = NEXT_SEQUENCE.getAndIncrement();
       EXECUTOR.execute(
           () -> {
             try {
-              handle(data);
+              handle(sequence, data);
             } catch (Throwable t) {
               LOG.log(Level.WARNING, "Unable to process: " + data, t);
               throw new DevToolsException(t);
@@ -242,7 +245,7 @@ public void onText(CharSequence data) {
     }
   }
 
-  private void handle(CharSequence data) {
+  private void handle(long sequence, CharSequence data) {
     // It's kind of gross to decode the data twice, but this lets us get started on something
     // that feels nice to users.
     // TODO: decode once, and once only
@@ -335,14 +338,14 @@ private void handle(CharSequence data) {
                       return;
                     }
 
-                    for (Consumer<?> action : event.getValue()) {
+                    for (BiConsumer<Long, ?> action : event.getValue()) {
                       @SuppressWarnings("unchecked")
-                      Consumer<Object> obj = (Consumer<Object>) action;
+                      BiConsumer<Long, Object> obj = (BiConsumer<Long, Object>) action;
                       LOG.log(
                           getDebugLogLevel(),
                           "Calling callback for {0} using {1} being passed {2}",
                           new Object[] {event.getKey(), obj, params});
-                      obj.accept(params);
+                      obj.accept(sequence, params);
                     }
                   }
                 });
diff --git a/java/src/org/openqa/selenium/devtools/DevTools.java b/java/src/org/openqa/selenium/devtools/DevTools.java
index f6cb7ad57edf3..b83e68f0fd8f0 100644
--- a/java/src/org/openqa/selenium/devtools/DevTools.java
+++ b/java/src/org/openqa/selenium/devtools/DevTools.java
@@ -26,6 +26,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeoutException;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.logging.Level;
@@ -89,10 +90,33 @@ public <X> X send(Command<X> command) {
     return connection.sendAndWait(cdpSession, command, timeout);
   }
 
+  /**
+   * Register a handler to the given event.
+   *
+   * @param event the event to listen to
+   * @param handler the handler to register
+   * @param <X> type of the data generated by the event
+   */
   public <X> void addListener(Event<X> event, Consumer<X> handler) {
     Require.nonNull("Event to listen for", event);
     Require.nonNull("Handler to call", handler);
 
+    connection.addListener(event, (sequence, x) -> handler.accept(x));
+  }
+
+  /**
+   * Register a handler to the given event, this handler will receive a sequence number to
+   * determinate the order of the data generated by the event. The sequence number might have gaps
+   * when other events are raised in between.
+   *
+   * @param event the event to listen to
+   * @param handler the handler to register
+   * @param <X> type of the data generated by the event
+   */
+  public <X> void addListener(Event<X> event, BiConsumer<Long, X> handler) {
+    Require.nonNull("Event to listen for", event);
+    Require.nonNull("Handler to call", handler);
+
     connection.addListener(event, handler);
   }