WebSockets Next: cancel returned Multi if the connection is closed

- fixes quarkusio#41025
mkouba · Jun 7, 2024 · 2af63ef · 2af63ef
1 parent 9222682
commit 2af63ef
Showing 1 changed file with 16 additions and 26 deletions.
diff --git a/...-next/runtime/src/main/java/io/quarkus/websockets/next/runtime/WebSocketEndpointBase.java b/...-next/runtime/src/main/java/io/quarkus/websockets/next/runtime/WebSocketEndpointBase.java
@@ -319,41 +319,31 @@ public Uni<Void> sendText(String message, boolean broadcast) {
         return broadcast ? connection.broadcast().sendText(message) : connection.sendText(message);
     }
 
-    public Uni<Void> multiText(Multi<Object> multi, Function<Object, Uni<Void>> action) {
-        multi.onFailure().recoverWithMulti(t -> doOnError(t).toMulti())
+    public Uni<Void> multiText(Multi<Object> multi, Function<? super Object, Uni<?>> action) {
+        multi
+                // Encode and send message
+                .onItem().call(action)
+                .onFailure().recoverWithMulti(t -> {
+                    return doOnError(t).toMulti();
+                })
                 .subscribe().with(
-                        m -> {
-                            // Encode and send message
-                            action.apply(m)
-                                    .onFailure().recoverWithUni(this::doOnError)
-                                    .subscribe()
-                                    .with(v -> LOG.debugf("Multi >> text message: %s", connection),
-                                            t -> LOG.errorf(t, "Unable to send text message from Multi: %s", connection));
-                        },
-                        t -> {
-                            LOG.errorf(t, "Unable to send text message from Multi: %s ", connection);
-                        });
+                        m -> LOG.debugf("Multi >> text message: %s", connection),
+                        t -> LOG.errorf(t, "Unable to send text message from Multi: %s ", connection));
         return Uni.createFrom().voidItem();
     }
 
     public Uni<Void> sendBinary(Buffer message, boolean broadcast) {
         return broadcast ? connection.broadcast().sendBinary(message) : connection.sendBinary(message);
     }
 
-    public Uni<Void> multiBinary(Multi<Object> multi, Function<Object, Uni<Void>> action) {
-        multi.onFailure().recoverWithMulti(t -> doOnError(t).toMulti())
+    public Uni<Void> multiBinary(Multi<Object> multi, Function<? super Object, Uni<?>> action) {
+        multi
+                // Encode and send message
+                .onItem().call(action)
+                .onFailure().recoverWithMulti(t -> doOnError(t).toMulti())
                 .subscribe().with(
-                        m -> {
-                            // Encode and send message
-                            action.apply(m)
-                                    .onFailure().recoverWithUni(this::doOnError)
-                                    .subscribe()
-                                    .with(v -> LOG.debugf("Multi >> binary message: %s", connection),
-                                            t -> LOG.errorf(t, "Unable to send binary message from Multi: %s", connection));
-                        },
-                        t -> {
-                            LOG.errorf(t, "Unable to send text message from Multi: %s ", connection);
-                        });
+                        m -> LOG.debugf("Multi >> binary message: %s", connection),
+                        t -> LOG.errorf(t, "Unable to send binary message from Multi: %s ", connection));
         return Uni.createFrom().voidItem();
     }
 }