Netty WebSocket Epoll: remove connect callback from handshake listener

Author: jknack

modules/jooby-netty/src/main/java/io/jooby/internal/netty/NettyContext.java

@@ -311,13 +311,8 @@ public NettyContext(ChannelHandlerContext ctx, HttpRequest req, Router router, S
       WebSocketServerHandshakerFactory factory = new WebSocketServerHandshakerFactory(webSocketURL,
           null, config);
       WebSocketServerHandshaker handshaker = factory.newHandshaker(webSocketRequest);
-      handshaker.handshake(ctx.channel(), webSocketRequest).addListener(future -> {
-        if (future.isSuccess()) {
-          webSocket.fireConnect();
-        } else {
-          handshaker.close(ctx.channel(), new CloseWebSocketFrame());
-        }
-      });
+      handshaker.handshake(ctx.channel(), webSocketRequest);
+      webSocket.fireConnect();
       Config conf = getRouter().getConfig();
       long timeout = conf.hasPath("websocket.idleTimeout")
           ? conf.getDuration("websocket.idleTimeout", TimeUnit.MINUTES)

modules/jooby-netty/src/main/java/io/jooby/internal/netty/NettyWebSocket.java

@@ -33,8 +33,6 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 public class NettyWebSocket implements WebSocketConfigurer, WebSocket, ChannelFutureListener {
   /** All connected websocket. */
@@ -168,7 +166,7 @@ private void handleMessage(WebSocketFrame frame) {
         }
         WebSocketMessage message = WebSocketMessage.create(getContext(), array(content));
 
-        fireCallback(webSocketTask(() -> messageCallback.onMessage(this, message)));
+        fireCallback(webSocketTask(() -> messageCallback.onMessage(this, message), false));
       } else {
         buffer = Unpooled.copiedBuffer(frame.content());
       }
@@ -177,21 +175,13 @@ private void handleMessage(WebSocketFrame frame) {
     }
   }
 
-  private void waitForConnect() {
-    try {
-      ready.await();
-    } catch (InterruptedException x) {
-      Thread.currentThread().interrupt();
-    }
-  }
-
   private void handleClose(WebSocketCloseStatus closeStatus) {
     try {
       if (isOpen()) {
         if (onCloseCallback != null) {
           Runnable closeCallback = () -> {
             try {
-              webSocketTask(() -> onCloseCallback.onClose(this, closeStatus)).run();
+              webSocketTask(() -> onCloseCallback.onClose(this, closeStatus), false).run();
             } finally {
               netty.ctx.channel()
                   .writeAndFlush(
@@ -228,31 +218,16 @@ private void handleError(Throwable x) {
   }
 
   void fireConnect() {
-    // fire only once
     addSession(this);
     if (connectCallback != null) {
       fireCallback(webSocketTask(() -> {
-        try {
-          connectCallback.onConnect(this);
-        } finally {
-          ready.countDown();
-        }
-      }));
+        connectCallback.onConnect(this);
+      }, true));
     } else {
       ready.countDown();
     }
   }
 
-  private Runnable webSocketTask(Runnable runnable) {
-    return () -> {
-      try {
-        runnable.run();
-      } catch (Throwable x) {
-        handleError(x);
-      }
-    };
-  }
-
   private void fireCallback(Runnable task) {
     if (dispatch) {
       Router router = netty.getRouter();
@@ -300,4 +275,25 @@ private void removeSession(NettyWebSocket ws) {
     }
   }
 
+  private Runnable webSocketTask(Runnable runnable, boolean isInit) {
+    return () -> {
+      try {
+        runnable.run();
+      } catch (Throwable x) {
+        handleError(x);
+      } finally {
+        if (isInit) {
+          ready.countDown();
+        }
+      }
+    };
+  }
+
+  private void waitForConnect() {
+    try {
+      ready.await();
+    } catch (InterruptedException x) {
+      Thread.currentThread().interrupt();
+    }
+  }
 }

modules/jooby-utow/src/main/java/io/jooby/internal/utow/UtowWebSocket.java

@@ -182,7 +182,7 @@ void fireConnect() {
 
   private void waitForConnect() {
     try {
-      ready.await(1, TimeUnit.MINUTES);
+      ready.await();
     } catch (InterruptedException x) {
       Thread.currentThread().interrupt();
     }