Synchronize on the underlying websocket (#10976)

shs96c · web-flow · commit 44d398ae6563 · 2022-08-30T11:15:50.000+02:00
This means that we can use the JDK11 websocket safely
in a multi-threaded environment.
diff --git a/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpClient.java b/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpClient.java
@@ -21,6 +21,7 @@
 import org.openqa.selenium.Credentials;
 import org.openqa.selenium.TimeoutException;
 import org.openqa.selenium.UsernameAndPassword;
+import org.openqa.selenium.WebDriverException;
 import org.openqa.selenium.remote.http.BinaryMessage;
 import org.openqa.selenium.remote.http.ClientConfig;
 import org.openqa.selenium.remote.http.CloseMessage;
@@ -46,27 +47,34 @@
 import java.net.http.HttpResponse.BodyHandlers;
 import java.net.http.HttpTimeoutException;
 import java.nio.ByteBuffer;
+import java.time.Duration;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
 
 import static java.net.http.HttpClient.Redirect.ALWAYS;
 
 public class JdkHttpClient implements HttpClient {
   private final JdkHttpMessages messages;
   private final java.net.http.HttpClient client;
+  private final Duration readTimeout;
 
   JdkHttpClient(ClientConfig config) {
     Objects.requireNonNull(config, "Client config must be set");
 
     this.messages = new JdkHttpMessages(config);
+    this.readTimeout = config.readTimeout();
 
     java.net.http.HttpClient.Builder builder = java.net.http.HttpClient.newBuilder()
       .connectTimeout(config.connectionTimeout())
       .followRedirects(ALWAYS);
 
-
     Credentials credentials = config.credentials();
     if (credentials != null) {
       if (!(credentials instanceof UsernameAndPassword)) {
@@ -148,18 +156,40 @@ public void onError(java.net.http.WebSocket webSocket, Throwable error) {
     return new WebSocket() {
       @Override
       public WebSocket send(Message message) {
+        Supplier<CompletableFuture<java.net.http.WebSocket>> makeCall;
+
         if (message instanceof BinaryMessage) {
           BinaryMessage binaryMessage = (BinaryMessage) message;
-          underlyingSocket.sendBinary(ByteBuffer.wrap(binaryMessage.data()), true);
+          makeCall = () -> underlyingSocket.sendBinary(ByteBuffer.wrap(binaryMessage.data()), true);
         } else if (message instanceof TextMessage) {
           TextMessage textMessage = (TextMessage) message;
-          underlyingSocket.sendText(textMessage.text(), true);
+          makeCall = () -> underlyingSocket.sendText(textMessage.text(), true);
         } else if (message instanceof CloseMessage) {
           CloseMessage closeMessage = (CloseMessage) message;
-          underlyingSocket.sendClose(closeMessage.code(), closeMessage.reason());
+          makeCall = () -> underlyingSocket.sendClose(closeMessage.code(), closeMessage.reason());
         } else {
           throw new IllegalArgumentException("Unsupport message type: " + message);
         }
+
+        synchronized (underlyingSocket) {
+          CompletableFuture<java.net.http.WebSocket> future = makeCall.get();
+          try {
+            future.get(readTimeout.toMillis(), TimeUnit.MILLISECONDS);
+          } catch (CancellationException e) {
+            throw new WebDriverException(e.getMessage(), e);
+          } catch (ExecutionException e) {
+            Throwable cause = e.getCause();
+            if (cause == null) {
+              throw new WebDriverException(e);
+            }
+            throw new WebDriverException(cause);
+          } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new WebDriverException(e.getMessage());
+          } catch (java.util.concurrent.TimeoutException e) {
+            throw new TimeoutException(e);
+          }
+        }
         return this;
       }
 
