[bidi][java] Ensure closed socket connection does not prevent freeing…

… up resources
SeleniumHQ · Oct 4, 2023 · 4bed5fa · 4bed5fa
1 parent b75c00f
commit 4bed5fa
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git a/java/src/org/openqa/selenium/bidi/Connection.java b/java/src/org/openqa/selenium/bidi/Connection.java
@@ -36,6 +36,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReadWriteLock;
@@ -72,13 +73,15 @@ public class Connection implements Closeable {
   private final ReadWriteLock callbacksLock = new ReentrantReadWriteLock(true);
   private final Multimap<Event<?>, Consumer<?>> eventCallbacks = HashMultimap.create();
   private final HttpClient client;
+  private final AtomicBoolean underlyingSocketClosed;
 
   public Connection(HttpClient client, String url) {
     Require.nonNull("HTTP client", client);
     Require.nonNull("URL to connect to", url);
 
     this.client = client;
     socket = this.client.openSocket(new HttpRequest(GET, url), new Listener());
+    underlyingSocketClosed = new AtomicBoolean();
   }
 
   private static class NamedConsumer<X> implements Consumer<X> {
@@ -210,7 +213,13 @@ public void clearListeners() {
       List<String> events =
           eventCallbacks.keySet().stream().map(Event::getMethod).collect(Collectors.toList());
 
-      send(new Command<>("session.unsubscribe", ImmutableMap.of("events", events)));
+      // If WebDriver close() is called, it closes the session if it is the last browsing context.
+      // It also closes the WebSocket from the remote end.
+      // If we try to now send commands, depending on the underlying web socket implementation, it will throw errors.
+      // Ideally, such errors should not prevent freeing up resources.
+      if (!underlyingSocketClosed.get()) {
+        send(new Command<>("session.unsubscribe", ImmutableMap.of("events", events)));
+      }
 
       eventCallbacks.clear();
     } finally {
@@ -237,6 +246,12 @@ public void onText(CharSequence data) {
             }
           });
     }
+
+    @Override
+    public void onClose(int code, String reason) {
+      LOG.fine("BiDi connection websocket closed");
+      underlyingSocketClosed.set(true);
+    }
   }
 
   private void handle(CharSequence data) {

diff --git a/java/test/org/openqa/selenium/bidi/BiDiSessionCleanUpTest.java b/java/test/org/openqa/selenium/bidi/BiDiSessionCleanUpTest.java
@@ -22,7 +22,7 @@
 
 import java.util.Collections;
 import org.junit.jupiter.api.Test;
-import org.openqa.selenium.TimeoutException;
+import org.openqa.selenium.WebDriverException;
 import org.openqa.selenium.WindowType;
 import org.openqa.selenium.firefox.FirefoxDriver;
 import org.openqa.selenium.firefox.FirefoxOptions;
@@ -78,7 +78,7 @@ void shouldCloseBiDiSessionIfLastWindowIsClosed() {
     driver.close();
 
     // Closing the last top-level browsing context, closes the WebDriver and BiDi session
-    assertThatExceptionOfType(TimeoutException.class)
+    assertThatExceptionOfType(WebDriverException.class)
         .isThrownBy(
             () ->
                 biDi.send(