From 4bed5fa7e236d93269186a34d20f661395c9b69d Mon Sep 17 00:00:00 2001
From: Puja Jagani <puja.jagani93@gmail.com>
Date: Wed, 4 Oct 2023 17:08:20 +0530
Subject: [PATCH] [bidi][java] Ensure closed socket connection does not prevent
 freeing up resources

---
 .../org/openqa/selenium/bidi/Connection.java    | 17 ++++++++++++++++-
 .../selenium/bidi/BiDiSessionCleanUpTest.java   |  4 ++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/java/src/org/openqa/selenium/bidi/Connection.java b/java/src/org/openqa/selenium/bidi/Connection.java
index b4f46bbbe7a93..cf9eb81ec1980 100644
--- 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,6 +73,7 @@ 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);
@@ -79,6 +81,7 @@ public Connection(HttpClient client, String 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
index 85964a529751a..090767dd89c1c 100644
--- 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(