[java] JdkHttpClient - Close all websockets before shutting down the executor (#12035)

joerg1985 · diemol · web-flow · commit 5c891ce374dd · 2023-05-15T19:46:43.000+02:00
Close all websockets before shutting down the executor

Co-authored-by: Diego Molina &lt;diemol@users.noreply.github.com&gt;
diff --git a/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpClient.java b/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpClient.java
@@ -50,6 +50,7 @@
 import java.net.http.HttpTimeoutException;
 import java.nio.ByteBuffer;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.CancellationException;
@@ -67,7 +68,7 @@ public class JdkHttpClient implements HttpClient {
   public static final Logger LOG = Logger.getLogger(JdkHttpClient.class.getName());
   private final JdkHttpMessages messages;
   private java.net.http.HttpClient client;
-  private WebSocket websocket;
+  private final List<WebSocket> websockets;
   private final ExecutorService executorService;
   private final Duration readTimeout;
 
@@ -76,7 +77,7 @@ public class JdkHttpClient implements HttpClient {
 
     this.messages = new JdkHttpMessages(config);
     this.readTimeout = config.readTimeout();
-
+    websockets = new ArrayList<>();
     executorService = Executors.newCachedThreadPool();
 
     java.net.http.HttpClient.Builder builder = java.net.http.HttpClient.newBuilder()
@@ -202,7 +203,7 @@ public void onError(java.net.http.WebSocket webSocket, Throwable error) {
 
     java.net.http.WebSocket underlyingSocket = webSocketCompletableFuture.join();
 
-    this.websocket = new WebSocket() {
+    WebSocket websocket = new WebSocket() {
       @Override
       public WebSocket send(Message message) {
         Supplier<CompletableFuture<java.net.http.WebSocket>> makeCall;
@@ -259,14 +260,11 @@ public WebSocket send(Message message) {
       @Override
       public void close() {
         LOG.fine("Closing websocket");
-        synchronized (underlyingSocket) {
-          if (!underlyingSocket.isOutputClosed()) {
-            underlyingSocket.sendClose(1000, "WebDriver closing socket");
-          }
-        }
+        send(new CloseMessage(1000, "WebDriver closing socket"));
       }
     };
-    return this.websocket;
+    websockets.add(websocket);
+    return websocket;
   }
 
   private URI getWebSocketUri(HttpRequest request) {
@@ -347,11 +345,18 @@ public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
 
   @Override
   public void close() {
-    executorService.shutdownNow();
-    if (this.websocket != null) {
-      this.websocket.close();
+    if (this.client == null) {
+      return ;
     }
     this.client = null;
+    for (WebSocket websocket : websockets) {
+      try {
+        websocket.close();
+      } catch (Exception e) {
+        LOG.log(Level.WARNING, "failed to close the websocket: " + websocket, e);
+      }
+    }
+    executorService.shutdownNow();
   }
 
   @AutoService(HttpClient.Factory.class)
