Close the websocket on error when using undertow [IMMUTANT-543]

When the client goes away unexpectedly, that will be expressed as a call to the on-error callback, but not on-close (when using undertow directly). In-container, both get called. This makes undertow behave the same way as in-container.
projectodd · Mar 31, 2015 · 0adea53 · 0adea53
1 parent 9ce817f
commit 0adea53
Showing 1 changed file with 17 additions and 16 deletions.
diff --git a/...src/main/java/org/projectodd/wunderboss/web/async/websocket/UndertowWebsocketChannel.java b/...src/main/java/org/projectodd/wunderboss/web/async/websocket/UndertowWebsocketChannel.java
@@ -38,7 +38,6 @@ public UndertowWebsocketChannel(final OnOpen onOpen,
 
     @Override
     public UndertowEndpoint endpoint() {
-        final UndertowWebsocketChannel channel = this;
         return new UndertowEndpoint() {
             @Override
             public void onMessage(WebSocketChannel _, Object message) {
@@ -48,8 +47,8 @@ public void onMessage(WebSocketChannel _, Object message) {
             @Override
             public void onOpen(WebSocketChannel baseChannel,
                                WebSocketHttpExchange exchange) {
-                channel.setUnderlyingChannel(baseChannel);
-                channel.notifyOpen(exchange);
+                setUnderlyingChannel(baseChannel);
+                notifyOpen(exchange);
             }
 
             @Override
@@ -60,6 +59,7 @@ public void onClose(WebSocketChannel _, CloseMessage message) {
             @Override
             public void onError(WebSocketChannel _, Throwable error) {
                 notifyError(error);
+                maybeCloseOnError(error);
             }
         };
     }
@@ -86,20 +86,16 @@ public boolean send(final Object message,
         final WebSocketCallback<Void> callback = new WebSocketCallback<Void>() {
             @Override
             public void complete(WebSocketChannel channel, Void context) {
-                Exception ex = null;
                 if (shouldClose) {
-                    try {
-                        close();
-                    } catch (IOException e) {
-                        ex = e;
-                    }
+                    close();
                 }
-                notifyComplete(onComplete, ex);
+                notifyComplete(onComplete, null);
             }
 
             @Override
             public void onError(WebSocketChannel channel, Void context, Throwable throwable) {
                 notifyComplete(onComplete, throwable);
+                maybeCloseOnError(throwable);
             }
         };
         if (message instanceof String) {
@@ -118,15 +114,20 @@ public void onError(WebSocketChannel channel, Void context, Throwable throwable)
     }
 
     @Override
-    public void close() throws IOException {
+    public void close() {
         if (isOpen()) {
-            this.underlyingChannel.sendClose();
-            notifyClose(CloseMessage.NORMAL_CLOSURE, "");
+            try {
+                this.underlyingChannel.sendClose();
+            } catch (IOException _) {}
+        }
+        notifyClose(CloseMessage.NORMAL_CLOSURE, "");
+    }
+
+    protected void maybeCloseOnError(Throwable error) {
+        if (error instanceof IOException) {
+            close();
         }
     }
 
     private WebSocketChannel underlyingChannel;
 }
-
-
-