WebSockets Next: perform connection cleanup during app shutdown

Author: mkouba

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

@@ -155,7 +155,7 @@ public Uni<WebSocketClientConnection> connect() {
             throw new WebSocketClientException(e);
         }
 
-        Uni<WebSocket> websocket = Uni.createFrom().<WebSocket> emitter(e -> {
+        Uni<WebSocketOpen> websocketOpen = Uni.createFrom().<WebSocketOpen> emitter(e -> {
             // Create a new event loop context for each client, otherwise the current context is used
             // We want to avoid a situation where if multiple clients/connections are created in a row,
             // the same event loop is used and so writing/receiving messages is de-facto serialized
@@ -171,7 +171,7 @@ public void handle(Void event) {
                             @Override
                             public void handle(AsyncResult<WebSocket> r) {
                                 if (r.succeeded()) {
-                                    e.complete(r.result());
+                                    e.complete(new WebSocketOpen(newCleanupConsumer(c, context), r.result()));
                                 } else {
                                     e.fail(r.cause());
                                 }
@@ -183,14 +183,20 @@ public void handle(AsyncResult<WebSocket> r) {
                 }
             });
         });
-        return websocket.map(ws -> {
+        return websocketOpen.map(wsOpen -> {
+            WebSocket ws = wsOpen.websocket();
             String clientId = BasicWebSocketConnector.class.getName();
             TrafficLogger trafficLogger = TrafficLogger.forClient(config);
-            WebSocketClientConnectionImpl connection = new WebSocketClientConnectionImpl(clientId, ws,
+            WebSocketClientConnectionImpl connection = new WebSocketClientConnectionImpl(clientId,
+                    ws,
                     codecs,
                     pathParams,
                     serverEndpointUri,
-                    headers, trafficLogger, userData, null);
+                    headers,
+                    trafficLogger,
+                    userData,
+                    null,
+                    wsOpen.cleanup());
             if (trafficLogger != null) {
                 trafficLogger.connectionOpened(connection);
             }

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

@@ -2,20 +2,21 @@
 
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map.Entry;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.stream.Stream;
 
-import jakarta.annotation.PreDestroy;
 import jakarta.enterprise.event.Event;
 import jakarta.inject.Singleton;
 
 import org.jboss.logging.Logger;
 
 import io.quarkus.arc.Arc;
 import io.quarkus.arc.ArcContainer;
+import io.quarkus.runtime.Shutdown;
 import io.quarkus.websockets.next.Closed;
 import io.quarkus.websockets.next.Open;
 import io.quarkus.websockets.next.OpenClientConnections;
@@ -26,7 +27,7 @@ public class ClientConnectionManager implements OpenClientConnections {
 
     private static final Logger LOG = Logger.getLogger(ClientConnectionManager.class);
 
-    private final ConcurrentMap<String, Set<WebSocketClientConnection>> endpointToConnections = new ConcurrentHashMap<>();
+    private final ConcurrentMap<String, Set<WebSocketClientConnectionImpl>> endpointToConnections = new ConcurrentHashMap<>();
 
     private final List<ClientConnectionListener> listeners = new CopyOnWriteArrayList<>();
 
@@ -50,10 +51,11 @@ public Iterator<WebSocketClientConnection> iterator() {
 
     @Override
     public Stream<WebSocketClientConnection> stream() {
-        return endpointToConnections.values().stream().flatMap(Set::stream).filter(WebSocketClientConnection::isOpen);
+        return endpointToConnections.values().stream().flatMap(Set::stream).filter(WebSocketClientConnection::isOpen)
+                .map(WebSocketClientConnection.class::cast);
     }
 
-    void add(String endpoint, WebSocketClientConnection connection) {
+    void add(String endpoint, WebSocketClientConnectionImpl connection) {
         LOG.debugf("Add client connection: %s", connection);
         if (endpointToConnections.computeIfAbsent(endpoint, e -> ConcurrentHashMap.newKeySet()).add(connection)) {
             if (openEvent != null) {
@@ -72,9 +74,9 @@ void add(String endpoint, WebSocketClientConnection connection) {
         }
     }
 
-    void remove(String endpoint, WebSocketClientConnection connection) {
+    void remove(String endpoint, WebSocketClientConnectionImpl connection) {
         LOG.debugf("Remove client connection: %s", connection);
-        Set<WebSocketClientConnection> connections = endpointToConnections.get(endpoint);
+        Set<WebSocketClientConnectionImpl> connections = endpointToConnections.get(endpoint);
         if (connections != null) {
             if (connections.remove(connection)) {
                 if (closedEvent != null) {
@@ -99,8 +101,8 @@ void remove(String endpoint, WebSocketClientConnection connection) {
      * @param endpoint
      * @return the connections for the given client endpoint, never {@code null}
      */
-    public Set<WebSocketClientConnection> getConnections(String endpoint) {
-        Set<WebSocketClientConnection> ret = endpointToConnections.get(endpoint);
+    public Set<WebSocketClientConnectionImpl> getConnections(String endpoint) {
+        Set<WebSocketClientConnectionImpl> ret = endpointToConnections.get(endpoint);
         if (ret == null) {
             return Set.of();
         }
@@ -111,9 +113,19 @@ public void addListener(ClientConnectionListener listener) {
         this.listeners.add(listener);
     }
 
-    @PreDestroy
-    void destroy() {
-        endpointToConnections.clear();
+    @Shutdown
+    void cleanup() {
+        if (!endpointToConnections.isEmpty()) {
+            int sum = 0;
+            for (Entry<String, Set<WebSocketClientConnectionImpl>> e : endpointToConnections.entrySet()) {
+                for (WebSocketClientConnectionImpl c : e.getValue()) {
+                    c.cleanup();
+                    sum++;
+                }
+            }
+            LOG.debugf("Cleanup performed for %s connections", sum);
+            endpointToConnections.clear();
+        }
     }
 
     public interface ClientConnectionListener {

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

@@ -6,6 +6,7 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
+import java.util.function.Consumer;
 
 import io.quarkus.websockets.next.HandshakeRequest;
 import io.quarkus.websockets.next.WebSocketClientConnection;
@@ -19,13 +20,17 @@ class WebSocketClientConnectionImpl extends WebSocketConnectionBase implements W
 
     private final WebSocket webSocket;
 
+    private final Consumer<WebSocketClientConnection> cleanup;
+
     WebSocketClientConnectionImpl(String clientId, WebSocket webSocket, Codecs codecs,
             Map<String, String> pathParams, URI serverEndpointUri, Map<String, List<String>> headers,
-            TrafficLogger trafficLogger, Map<String, Object> userData, SendingInterceptor sendingInterceptor) {
+            TrafficLogger trafficLogger, Map<String, Object> userData, SendingInterceptor sendingInterceptor,
+            Consumer<WebSocketClientConnection> cleanup) {
         super(Map.copyOf(pathParams), codecs, new ClientHandshakeRequestImpl(serverEndpointUri, headers), trafficLogger,
                 new UserDataImpl(userData), sendingInterceptor);
         this.clientId = clientId;
         this.webSocket = Objects.requireNonNull(webSocket);
+        this.cleanup = cleanup;
     }
 
     @Override
@@ -48,6 +53,12 @@ public int hashCode() {
         return Objects.hash(identifier);
     }
 
+    protected void cleanup() {
+        if (cleanup != null) {
+            cleanup.accept(this);
+        }
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (this == obj)

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

@@ -11,23 +11,32 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Consumer;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+import org.jboss.logging.Logger;
+
 import io.quarkus.tls.TlsConfiguration;
 import io.quarkus.tls.TlsConfigurationRegistry;
 import io.quarkus.tls.runtime.config.TlsConfigUtils;
 import io.quarkus.websockets.next.UserData.TypedKey;
+import io.quarkus.websockets.next.WebSocketClientConnection;
 import io.quarkus.websockets.next.WebSocketClientException;
 import io.quarkus.websockets.next.runtime.config.WebSocketsClientRuntimeConfig;
 import io.vertx.core.Vertx;
+import io.vertx.core.http.WebSocket;
+import io.vertx.core.http.WebSocketClient;
 import io.vertx.core.http.WebSocketClientOptions;
 import io.vertx.core.http.WebSocketConnectOptions;
+import io.vertx.core.impl.ContextImpl;
 
 abstract class WebSocketConnectorBase<THIS extends WebSocketConnectorBase<THIS>> {
 
     protected static final Pattern PATH_PARAM_PATTERN = Pattern.compile("\\{[a-zA-Z0-9_]+\\}");
 
+    private static final Logger LOG = Logger.getLogger(WebSocketConnectorBase.class);
+
     // mutable state
 
     protected URI baseUri;
@@ -201,4 +210,28 @@ protected WebSocketConnectOptions newConnectOptions(URI serverEndpointUri) {
     protected boolean isSecure(URI uri) {
         return "https".equals(uri.getScheme()) || "wss".equals(uri.getScheme());
     }
+
+    record WebSocketOpen(Consumer<WebSocketClientConnection> cleanup, WebSocket websocket) {
+    }
+
+    Consumer<WebSocketClientConnection> newCleanupConsumer(WebSocketClient client, ContextImpl context) {
+        return new Consumer<WebSocketClientConnection>() {
+            @Override
+            public void accept(WebSocketClientConnection conn) {
+                try {
+                    client.close();
+                    LOG.debugf("Client closed for connection %s", conn.id());
+                } catch (Throwable e) {
+                    LOG.errorf(e, "Unable to close the client for connection %s", conn.id());
+                }
+                try {
+                    context.close();
+                    LOG.debugf("Context closed for connection %s", conn.id());
+                } catch (Throwable e) {
+                    LOG.errorf(e, "Unable to close the context for connection %s", conn.id());
+                }
+            }
+        };
+    }
+
 }

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

@@ -102,7 +102,7 @@ public Uni<WebSocketClientConnection> connect() {
 
         var telemetrySupport = telemetryProvider == null ? null
                 : telemetryProvider.createClientTelemetrySupport(clientEndpoint.path);
-        Uni<WebSocket> websocket = Uni.createFrom().<WebSocket> emitter(e -> {
+        Uni<WebSocketOpen> websocketOpen = Uni.createFrom().<WebSocketOpen> emitter(e -> {
             // Create a new event loop context for each client, otherwise the current context is used
             // We want to avoid a situation where if multiple clients/connections are created in a row,
             // the same event loop is used and so writing/receiving messages is de-facto serialized
@@ -121,7 +121,7 @@ public void handle(Void event) {
                             @Override
                             public void handle(AsyncResult<WebSocket> r) {
                                 if (r.succeeded()) {
-                                    e.complete(r.result());
+                                    e.complete(new WebSocketOpen(newCleanupConsumer(c, context), r.result()));
                                 } else {
                                     if (telemetrySupport != null && telemetrySupport.interceptConnection()) {
                                         telemetrySupport.connectionOpeningFailed(r.cause());
@@ -136,13 +136,20 @@ public void handle(AsyncResult<WebSocket> r) {
                 }
             });
         });
-        return websocket.map(ws -> {
+        return websocketOpen.map(wsOpen -> {
+            WebSocket ws = wsOpen.websocket();
             TrafficLogger trafficLogger = TrafficLogger.forClient(config);
             SendingInterceptor sendingInterceptor = telemetrySupport == null ? null : telemetrySupport.getSendingInterceptor();
-            WebSocketClientConnectionImpl connection = new WebSocketClientConnectionImpl(clientEndpoint.clientId, ws,
+            WebSocketClientConnectionImpl connection = new WebSocketClientConnectionImpl(clientEndpoint.clientId,
+                    ws,
                     codecs,
                     pathParams,
-                    serverEndpointUri, headers, trafficLogger, userData, sendingInterceptor);
+                    serverEndpointUri,
+                    headers,
+                    trafficLogger,
+                    userData,
+                    sendingInterceptor,
+                    wsOpen.cleanup());
             if (trafficLogger != null) {
                 trafficLogger.connectionOpened(connection);
             }