Merge pull request #50630 from michalvavrik/feature/ws-next-client-customization

WebSockets Next client: make it possible to customize the WebSocketConnectOptions and WebSocketClientOptions

Author: mkouba

docs/src/main/asciidoc/websockets-next-reference.adoc

@@ -1291,21 +1291,25 @@ public class MyBean {
         WebSocketClientConnection connection = connector
             .baseUri(uri) <2>
             .pathParam("name", "Roxanne") <3>
+            .customizeOptions((connectOptions, clientOptions) -> {  <4>
+               // ...
+            })
             .connectAndAwait();
-        connection.sendTextAndAwait("Hi!"); <4>
+        connection.sendTextAndAwait("Hi!"); <5>
     }
 }
 ----
 <1> Inject the connector for `ClientEndpoint`.
 <2> If the base `URI` is not supplied we attempt to obtain the value from the config. The key consists of the client id and the `.base-uri` suffix.
 <3> Set the path param value. Throws `IllegalArgumentException` if the client endpoint path does not contain a parameter with the given name.
-<4> Use the connection to send messages, if needed.
+<4> Optionally, customize client and connection options not configurable with the connector API or Quarkus configuration.
+<5> Use the connection to send messages, if needed.
 
 NOTE: If an application attempts to inject a connector for a missing endpoint, an error is thrown.
 
 Connectors are not thread-safe and should not be used concurrently.
 Connectors should also not be reused. 
-If you need to create multiple connections in a row you'll need to obtain a new connetor instance programmatically using `Instance#get()`:
+If you need to create multiple connections in a row you'll need to obtain a new connector instance programmatically using `Instance#get()`:
 
 [source, java]
 ----
@@ -1350,6 +1354,9 @@ public class MyBean {
             .onTextMessage((c, m) -> { <5>
                // ...
             })
+            .customizeOptions((connectOptions, clientOptions) -> {  <6>
+               // ...
+            })
             .connectAndAwait();
     }
 }
@@ -1359,6 +1366,7 @@ public class MyBean {
 <3> The additional path that should be appended to the base URI.
 <4> Set the execution model for callback handlers. By default, the callback may block the current thread. However in this case, the callback is executed on the event loop and may not block the current thread.
 <5> The lambda will be called for every text message sent from the server.
+<6> Optionally, customize client and connection options not configurable with the connector API or Quarkus configuration.
 
 The basic connector is closer to a low-level API and is reserved for advanced users.
 However, unlike others low-level WebSocket clients, it is still a CDI bean and can be injected in other beans.

extensions/websockets-next/deployment/src/test/java/io/quarkus/websockets/next/test/client/BasicConnectorContextTest.java

@@ -15,6 +15,7 @@
 import io.quarkus.test.QuarkusUnitTest;
 import io.quarkus.test.common.http.TestHTTPResource;
 import io.quarkus.websockets.next.BasicWebSocketConnector;
+import io.quarkus.websockets.next.HandshakeRequest;
 import io.quarkus.websockets.next.OnClose;
 import io.quarkus.websockets.next.OnOpen;
 import io.quarkus.websockets.next.WebSocket;
@@ -35,6 +36,8 @@ public class BasicConnectorContextTest {
 
     static final Set<String> THREADS = ConcurrentHashMap.newKeySet();
 
+    static final Set<String> MESSAGES = ConcurrentHashMap.newKeySet();
+
     static final CountDownLatch CLOSED_LATCH = new CountDownLatch(2);
 
     @Test
@@ -46,10 +49,14 @@ void testClient() throws InterruptedException {
                     String thread = Thread.currentThread().getName();
                     THREADS.add(thread);
                     MESSAGE_LATCH.countDown();
+                    MESSAGES.add(m);
                 })
                 .onClose((c, cr) -> {
                     CLOSED_LATCH.countDown();
                 })
+                .customizeOptions(((connectOptions, clientOptions) -> {
+                    connectOptions.addHeader("Foo", "Eric");
+                }))
                 .baseUri(uri);
         WebSocketClientConnection conn1 = connector.connectAndAwait();
         WebSocketClientConnection conn2 = connector.connectAndAwait();
@@ -67,6 +74,8 @@ void testClient() throws InterruptedException {
         conn2.closeAndAwait();
         assertTrue(ServerEndpoint.CLOSED_LATCH.await(5, TimeUnit.SECONDS));
         assertTrue(CLOSED_LATCH.await(5, TimeUnit.SECONDS));
+        assertTrue(MESSAGES.stream().allMatch("Hello Eric!"::equals),
+                () -> "Expected messages equal to 'Hello Eric!', but got " + MESSAGES);
     }
 
     @WebSocket(path = "/end")
@@ -75,8 +84,8 @@ public static class ServerEndpoint {
         static final CountDownLatch CLOSED_LATCH = new CountDownLatch(1);
 
         @OnOpen
-        String open() {
-            return "Hello!";
+        String open(HandshakeRequest handshakeRequest) {
+            return "Hello " + handshakeRequest.header("Foo") + "!";
         }
 
         @OnClose

extensions/websockets-next/deployment/src/test/java/io/quarkus/websockets/next/test/client/BasicConnectorTest.java

@@ -1,18 +1,22 @@
 package io.quarkus.websockets.next.test.client;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.net.URI;
+import java.net.UnknownHostException;
+import java.time.Duration;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
 
 import jakarta.inject.Inject;
 
@@ -32,6 +36,9 @@
 import io.quarkus.websockets.next.WebSocket;
 import io.quarkus.websockets.next.WebSocketClientConnection;
 import io.vertx.core.Context;
+import io.vertx.core.http.WebSocketClientOptions;
+import io.vertx.core.http.WebSocketConnectOptions;
+import io.vertx.core.net.ProxyOptions;
 
 public class BasicConnectorTest {
 
@@ -117,12 +124,35 @@ void testClient() throws InterruptedException {
         assertTrue(ServerEndpoint.CLOSED_LATCH.await(5, TimeUnit.SECONDS));
 
         CountDownLatch conn2Latch = new CountDownLatch(1);
-        WebSocketClientConnection connection2 = BasicWebSocketConnector
+        WebSocketClientConnection connection2 = createConnection2(conn2Latch, (o1, o2) -> {
+        });
+        assertNotNull(connection2);
+        assertTrue(connection2.handshakeRequest().localAddress().matches("127\\.0\\.0\\.1:\\d+"));
+        assertEquals("127.0.0.1:8081", connection2.handshakeRequest().remoteAddress());
+        assertTrue(conn2Latch.await(5, TimeUnit.SECONDS));
+    }
+
+    @Test
+    void testClientOptionsCustomization() {
+        // same connection with empty customizer succeeds in the 'testClient' test method,
+        // together with the thrown exception, this verifies that client options customization works
+        assertThatThrownBy(() -> createConnection2(new CountDownLatch(1), (ignored, clientOptions) -> {
+            clientOptions.setProxyOptions(new ProxyOptions()
+                    .setHost("robert")
+                    .setPort(999)
+                    .setConnectTimeout(Duration.ofMillis(500)));
+        })).rootCause().isInstanceOf(UnknownHostException.class).hasMessageContaining("robert");
+    }
+
+    private WebSocketClientConnection createConnection2(CountDownLatch conn2Latch,
+            BiConsumer<WebSocketConnectOptions, WebSocketClientOptions> customizer) {
+        return BasicWebSocketConnector
                 .create()
                 .baseUri(uri)
                 .path("/Cool")
                 .executionModel(ExecutionModel.NON_BLOCKING)
                 .addHeader("X-Test", "foo")
+                .customizeOptions(customizer)
                 .onTextMessage((c, m) -> {
                     assertTrue(Context.isOnEventLoopThread());
                     // Path params not set
@@ -132,10 +162,6 @@ void testClient() throws InterruptedException {
                     conn2Latch.countDown();
                 })
                 .connectAndAwait();
-        assertNotNull(connection2);
-        assertTrue(connection2.handshakeRequest().localAddress().matches("127\\.0\\.0\\.1:\\d+"));
-        assertEquals("127.0.0.1:8081", connection2.handshakeRequest().remoteAddress());
-        assertTrue(conn2Latch.await(5, TimeUnit.SECONDS));
     }
 
     @WebSocket(path = "/end/{name}")

extensions/websockets-next/deployment/src/test/java/io/quarkus/websockets/next/test/client/ClientEndpointTest.java

@@ -1,10 +1,13 @@
 package io.quarkus.websockets.next.test.client;
 
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.net.URI;
+import java.net.UnknownHostException;
+import java.time.Duration;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
@@ -30,6 +33,7 @@
 import io.quarkus.websockets.next.WebSocketConnection;
 import io.quarkus.websockets.next.WebSocketConnector;
 import io.vertx.core.buffer.Buffer;
+import io.vertx.core.net.ProxyOptions;
 
 public class ClientEndpointTest {
 
@@ -68,6 +72,19 @@ void testClient() throws InterruptedException {
         connection.closeAndAwait();
         assertTrue(ClientEndpoint.CLOSED_LATCH.await(5, TimeUnit.SECONDS));
         assertTrue(ServerEndpoint.CLOSED_LATCH.await(5, TimeUnit.SECONDS));
+
+        // use the same connector instance, but this time configure unknown host and expect failure
+        assertThatThrownBy(() -> connector
+                .baseUri(uri)
+                .pathParam("name", "Lu=")
+                .customizeOptions((ignored, clientOptions) -> {
+                    clientOptions.setProxyOptions(new ProxyOptions()
+                            .setHost("robert")
+                            .setPort(999)
+                            .setConnectTimeout(Duration.ofMillis(500)));
+                })
+                .connectAndAwait())
+                .rootCause().isInstanceOf(UnknownHostException.class).hasMessageContaining("robert");
     }
 
     @WebSocket(path = "/endpoint/{name}")

extensions/websockets-next/deployment/src/test/java/io/quarkus/websockets/next/test/client/programmatic/ClientEndpointProgrammaticTest.java

@@ -55,6 +55,13 @@ void testClient() throws InterruptedException {
                 .userData(TypedKey.forInt("int"), Integer.MAX_VALUE)
                 .userData(TypedKey.forLong("long"), Long.MAX_VALUE)
                 .userData(TypedKey.forString("string"), "Lu")
+                .customizeOptions((connectOptions, clientOptions) -> {
+                    connectOptions
+                            // test that port configuration is overridden
+                            .setPort(123456)
+                            // test that "Bar" header is sent with the opening handshake request
+                            .addHeader("Bar", "Adam");
+                })
                 .connectAndAwait();
         assertTrue(connection1.userData().get(TypedKey.forBoolean("boolean")));
         assertEquals(Integer.MAX_VALUE, connection1.userData().get(TypedKey.forInt("int")));
@@ -89,7 +96,7 @@ void testClient() throws InterruptedException {
         assertEquals("Ma", ClientEndpoint.CONNECTION_USER_DATA.get(connection2.id()).get(TypedKey.forString("string")));
 
         assertTrue(ClientEndpoint.MESSAGE_LATCH.await(5, TimeUnit.SECONDS));
-        assertTrue(ClientEndpoint.MESSAGES.contains("Lu:Hello Lu!"));
+        assertTrue(ClientEndpoint.MESSAGES.contains("Lu:Hello Lu Adam!"));
         assertTrue(ClientEndpoint.MESSAGES.contains("Lu:Hi!"));
         assertTrue(ClientEndpoint.MESSAGES.contains("Ma:Hello Ma!"));
         assertTrue(ClientEndpoint.MESSAGES.contains("Ma:Hi!"), ClientEndpoint.MESSAGES.toString());
@@ -107,7 +114,12 @@ public static class ServerEndpoint {
 
         @OnOpen
         String open(HandshakeRequest handshakeRequest) {
-            return "Hello " + handshakeRequest.header("Foo") + "!";
+            StringBuilder result = new StringBuilder("Hello " + handshakeRequest.header("Foo"));
+            if (handshakeRequest.header("Bar") != null) {
+                result.append(" ").append(handshakeRequest.header("Bar"));
+            }
+            result.append("!");
+            return result.toString();
         }
 
         @OnTextMessage

extensions/websockets-next/runtime/src/main/java/io/quarkus/websockets/next/BasicWebSocketConnector.java

@@ -14,6 +14,8 @@
 import io.smallrye.common.annotation.CheckReturnValue;
 import io.smallrye.mutiny.Uni;
 import io.vertx.core.buffer.Buffer;
+import io.vertx.core.http.WebSocketClientOptions;
+import io.vertx.core.http.WebSocketConnectOptions;
 
 /**
  * A basic connector can be used to configure and open a new client connection. Unlike with {@link WebSocketConnector} a
@@ -245,4 +247,15 @@ enum ExecutionModel {
         VIRTUAL_THREAD,
     }
 
+    /**
+     * This method allows to customize {@link WebSocketConnectOptions} and {@link WebSocketClientOptions}.
+     * Connection options configured directly with this API, such as host and port, have higher priority than this customizer.
+     * Client options configured directly with Quarkus configuration (e.g. if you configure the maximum size of a frame)
+     * have also higher priority than this customizer.
+     * Purpose of this customizer is to complement configuration options not configured elsewhere.
+     *
+     * @param customizer options customizer; must not be null
+     * @return self
+     */
+    BasicWebSocketConnector customizeOptions(BiConsumer<WebSocketConnectOptions, WebSocketClientOptions> customizer);
 }

extensions/websockets-next/runtime/src/main/java/io/quarkus/websockets/next/WebSocketConnector.java

@@ -2,6 +2,7 @@
 
 import java.net.URI;
 import java.net.URLEncoder;
+import java.util.function.BiConsumer;
 
 import jakarta.enterprise.inject.Default;
 import jakarta.enterprise.inject.Instance;
@@ -10,6 +11,8 @@
 import io.quarkus.websockets.next.UserData.TypedKey;
 import io.smallrye.common.annotation.CheckReturnValue;
 import io.smallrye.mutiny.Uni;
+import io.vertx.core.http.WebSocketClientOptions;
+import io.vertx.core.http.WebSocketConnectOptions;
 
 /**
  * A connector can be used to configure and open a new client connection backed by a client endpoint that is used to
@@ -132,4 +135,16 @@ default WebSocketClientConnection connectAndAwait() {
         return connect().await().indefinitely();
     }
 
+    /**
+     * This method allows to customize {@link WebSocketConnectOptions} and {@link WebSocketClientOptions}.
+     * Connection options configured directly with this API, such as host and port, have higher priority than this customizer.
+     * Client options configured directly with Quarkus configuration (e.g. if you configure the maximum size of a frame)
+     * have also higher priority than this customizer.
+     * Purpose of this customizer is to complement configuration options not configured elsewhere.
+     *
+     * @param customizer options customizer; must not be null
+     * @return self
+     */
+    WebSocketConnector<CLIENT> customizeOptions(BiConsumer<WebSocketConnectOptions, WebSocketClientOptions> customizer);
+
 }

extensions/websockets-next/runtime/src/main/java/io/quarkus/websockets/next/runtime/WebSocketConnectorBase.java

@@ -13,6 +13,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -43,6 +44,10 @@ abstract class WebSocketConnectorBase<THIS extends WebSocketConnectorBase<THIS>>
 
     protected URI baseUri;
 
+    protected WebSocketConnectOptions customWebSocketConnectOptions;
+
+    protected WebSocketClientOptions customWebSocketClientOptions;
+
     protected final Map<String, String> pathParams;
 
     protected final Map<String, List<String>> headers;
@@ -83,6 +88,8 @@ abstract class WebSocketConnectorBase<THIS extends WebSocketConnectorBase<THIS>>
         this.path = "";
         this.pathParamNames = Set.of();
         this.userData = new HashMap<>();
+        this.customWebSocketConnectOptions = null;
+        this.customWebSocketClientOptions = null;
     }
 
     public THIS baseUri(URI baseUri) {
@@ -170,7 +177,12 @@ String replacePathParameters(String path) {
     }
 
     protected WebSocketClientOptions populateClientOptions() {
-        WebSocketClientOptions clientOptions = new WebSocketClientOptions();
+        final WebSocketClientOptions clientOptions;
+        if (customWebSocketClientOptions != null) {
+            clientOptions = new WebSocketClientOptions(customWebSocketClientOptions);
+        } else {
+            clientOptions = new WebSocketClientOptions();
+        }
         if (config.offerPerMessageCompression()) {
             clientOptions.setTryUsePerMessageCompression(true);
             if (config.compressionLevel().isPresent()) {
@@ -210,7 +222,13 @@ protected WebSocketClientOptions populateClientOptions() {
     }
 
     protected WebSocketConnectOptions newConnectOptions(URI serverEndpointUri) {
-        WebSocketConnectOptions connectOptions = new WebSocketConnectOptions()
+        final WebSocketConnectOptions connectOptions;
+        if (customWebSocketConnectOptions != null) {
+            connectOptions = new WebSocketConnectOptions(customWebSocketConnectOptions);
+        } else {
+            connectOptions = new WebSocketConnectOptions();
+        }
+        connectOptions
                 .setSsl(isSecure(serverEndpointUri))
                 .setHost(serverEndpointUri.getHost());
         if (serverEndpointUri.getPort() != -1) {
@@ -249,4 +267,15 @@ public void accept(WebSocketClientConnection conn) {
         };
     }
 
+    public THIS customizeOptions(BiConsumer<WebSocketConnectOptions, WebSocketClientOptions> customizer) {
+        Objects.requireNonNull(customizer, "Options customizer must not be null");
+        if (customWebSocketClientOptions == null) {
+            customWebSocketClientOptions = new WebSocketClientOptions();
+        }
+        if (customWebSocketConnectOptions == null) {
+            customWebSocketConnectOptions = new WebSocketConnectOptions();
+        }
+        customizer.accept(customWebSocketConnectOptions, customWebSocketClientOptions);
+        return self();
+    }
 }