From 50a1b31a8ae8f9c4b6ef455bf2ae66ac6ef26683 Mon Sep 17 00:00:00 2001
From: Joakim Erdfelt <joakim.erdfelt@gmail.com>
Date: Fri, 22 Sep 2023 05:08:50 -0700
Subject: [PATCH] Issue #10547 - Allow Executor of WebSocketClient to be
 customized via HttpClient (#10548)

* Issue #10547 - Allow Executor of WebSocketClient to be customized via HttpClient

Signed-off-by: Joakim Erdfelt <joakim.erdfelt@gmail.com>
---
 .../websocket/client/WebSocketClient.java     | 41 +++++++-
 .../tests/client/WebSocketClientTest.java     | 93 +++++++++++++++++++
 2 files changed, 130 insertions(+), 4 deletions(-)

diff --git a/jetty-websocket/websocket-jetty-client/src/main/java/org/eclipse/jetty/websocket/client/WebSocketClient.java b/jetty-websocket/websocket-jetty-client/src/main/java/org/eclipse/jetty/websocket/client/WebSocketClient.java
index 27195651093..d60682736a6 100644
--- a/jetty-websocket/websocket-jetty-client/src/main/java/org/eclipse/jetty/websocket/client/WebSocketClient.java
+++ b/jetty-websocket/websocket-jetty-client/src/main/java/org/eclipse/jetty/websocket/client/WebSocketClient.java
@@ -64,7 +64,7 @@ public class WebSocketClient extends ContainerLifeCycle implements WebSocketPoli
     private final List<WebSocketSessionListener> sessionListeners = new CopyOnWriteArrayList<>();
     private final SessionTracker sessionTracker = new SessionTracker();
     private final Configuration.ConfigurationCustomizer configurationCustomizer = new Configuration.ConfigurationCustomizer();
-    private final WebSocketComponents components = new WebSocketComponents();
+    private final WebSocketComponents components;
     private boolean stopAtShutdown = false;
     private long _stopTimeout = Long.MAX_VALUE;
 
@@ -77,12 +77,45 @@ public class WebSocketClient extends ContainerLifeCycle implements WebSocketPoli
     }
 
     /**
-     * Instantiate a WebSocketClient using HttpClient for defaults
+     * <p>
+     * Instantiate a WebSocketClient.
+     * </p>
+     *
+     * <p>
+     *     HTTP behaviors of the WebSocket upgrade are taken from the HttpClient configuration.
+     * </p>
      *
      * @param httpClient the HttpClient to base internal defaults off of
      */
     public WebSocketClient(HttpClient httpClient)
     {
+        this (httpClient,
+            new WebSocketComponents(
+                null,
+                null,
+                httpClient != null ? httpClient.getByteBufferPool() : null,
+                null,
+                null,
+                httpClient != null ? httpClient.getExecutor() : null)
+        );
+    }
+
+    /**
+     * <p>
+     * Instantiate a WebSocketClient.
+     * </p>
+     *
+     * <p>
+     *     HTTP behaviors of the WebSocket upgrade are taken from the {@link HttpClient} configuration.
+     *     WebSocket behaviors are taken from the {@link WebSocketComponents} configuration.
+     * </p>
+     *
+     * @param httpClient the HttpClient to use for the HTTP behaviors of WebSocket upgrade
+     * @param webSocketComponents the WebSocketComponents to use for WebSocket behaviors
+     */
+    public WebSocketClient(HttpClient httpClient, WebSocketComponents webSocketComponents)
+    {
+        components = webSocketComponents;
         coreClient = new WebSocketCoreClient(httpClient, components);
         addManaged(coreClient);
         frameHandlerFactory = new JettyWebSocketFrameHandlerFactory(this, components);
@@ -336,13 +369,13 @@ public class WebSocketClient extends ContainerLifeCycle implements WebSocketPoli
 
     public ByteBufferPool getBufferPool()
     {
-        return getHttpClient().getByteBufferPool();
+        return components.getBufferPool();
     }
 
     @Override
     public Executor getExecutor()
     {
-        return getHttpClient().getExecutor();
+        return components.getExecutor();
     }
 
     public HttpClient getHttpClient()
diff --git a/jetty-websocket/websocket-jetty-tests/src/test/java/org/eclipse/jetty/websocket/tests/client/WebSocketClientTest.java b/jetty-websocket/websocket-jetty-tests/src/test/java/org/eclipse/jetty/websocket/tests/client/WebSocketClientTest.java
index 8756812bfb6..0fd9ad410b4 100644
--- a/jetty-websocket/websocket-jetty-tests/src/test/java/org/eclipse/jetty/websocket/tests/client/WebSocketClientTest.java
+++ b/jetty-websocket/websocket-jetty-tests/src/test/java/org/eclipse/jetty/websocket/tests/client/WebSocketClientTest.java
@@ -22,9 +22,15 @@ import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
+import org.eclipse.jetty.client.HttpClient;
+import org.eclipse.jetty.client.HttpClientTransport;
+import org.eclipse.jetty.client.http.HttpClientTransportOverHTTP;
+import org.eclipse.jetty.io.ClientConnector;
 import org.eclipse.jetty.server.Server;
 import org.eclipse.jetty.server.ServerConnector;
 import org.eclipse.jetty.servlet.ServletContextHandler;
@@ -36,6 +42,7 @@ import org.eclipse.jetty.websocket.api.UpgradeRequest;
 import org.eclipse.jetty.websocket.api.util.WSURI;
 import org.eclipse.jetty.websocket.client.ClientUpgradeRequest;
 import org.eclipse.jetty.websocket.client.WebSocketClient;
+import org.eclipse.jetty.websocket.core.WebSocketComponents;
 import org.eclipse.jetty.websocket.server.config.JettyWebSocketServletContainerInitializer;
 import org.eclipse.jetty.websocket.tests.AnnoMaxMessageEndpoint;
 import org.eclipse.jetty.websocket.tests.CloseTrackingEndpoint;
@@ -55,6 +62,7 @@ import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.hamcrest.Matchers.nullValue;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class WebSocketClientTest
@@ -112,6 +120,91 @@ public class WebSocketClientTest
         server.stop();
     }
 
+    @Test
+    public void testCustomizeExecutorDirectly() throws Exception
+    {
+        Executor executor = Executors.newFixedThreadPool(50);
+        HttpClient httpClient = new HttpClient();
+        httpClient.setExecutor(executor);
+        try
+        {
+            httpClient.start();
+            WebSocketClient webSocketClient = new WebSocketClient(httpClient);
+            try
+            {
+                webSocketClient.start();
+                Executor inuseExecutor = webSocketClient.getExecutor();
+                assertSame(executor, inuseExecutor);
+            }
+            finally
+            {
+                webSocketClient.stop();
+            }
+        }
+        finally
+        {
+            httpClient.stop();
+        }
+    }
+
+    @Test
+    public void testCustomizeWebSocketComponentsExecutor() throws Exception
+    {
+        HttpClient httpClient = new HttpClient();
+        try
+        {
+            httpClient.start();
+            Executor executor = Executors.newFixedThreadPool(50);
+            WebSocketComponents webSocketComponents = new WebSocketComponents(null, null,
+                null, null, null, executor);
+            WebSocketClient webSocketClient = new WebSocketClient(httpClient, webSocketComponents);
+            try
+            {
+                webSocketClient.start();
+                Executor inuseExecutor = webSocketClient.getExecutor();
+                assertSame(executor, inuseExecutor);
+            }
+            finally
+            {
+                webSocketClient.stop();
+            }
+        }
+        finally
+        {
+            httpClient.stop();
+        }
+    }
+
+    @Test
+    public void testCustomizeExecutorViaConnector() throws Exception
+    {
+        ClientConnector clientConnector = new ClientConnector();
+        clientConnector.setSelectors(1);
+        Executor executor = Executors.newFixedThreadPool(50);
+        clientConnector.setExecutor(executor);
+        HttpClientTransport transport = new HttpClientTransportOverHTTP(clientConnector);
+        HttpClient httpClient = new HttpClient(transport);
+        try
+        {
+            httpClient.start();
+            WebSocketClient webSocketClient = new WebSocketClient(httpClient);
+            try
+            {
+                webSocketClient.start();
+                Executor inuseExecutor = webSocketClient.getExecutor();
+                assertSame(executor, inuseExecutor);
+            }
+            finally
+            {
+                webSocketClient.stop();
+            }
+        }
+        finally
+        {
+            httpClient.stop();
+        }
+    }
+
     @Test
     public void testAddExtensionNotInstalled() throws Exception
     {