Clean up buffered WebSocketFrames handling

slandelle · slandelle · commit 385578abc293 · 2017-03-08T22:35:05.000+01:00
Motivation:

* WebSocketUpgradeHandler should be a AsyncHandler&lt;NettyWebSocket&gt;, no
need to hide implementation there and cast all over the place!
* buffer frames in the WebSocket, not in the WebSocketUpgradeHandler
* Release pending buffered frames when connection is abruptly closed or
crashes

Modifications:

See above

Result:

* cleaner code
* no more memory leak on WebSocket abrupt close or crash
diff --git a/client/src/main/java/org/asynchttpclient/netty/handler/WebSocketHandler.java b/client/src/main/java/org/asynchttpclient/netty/handler/WebSocketHandler.java
@@ -22,11 +22,6 @@
 import io.netty.handler.codec.http.HttpRequest;
 import io.netty.handler.codec.http.HttpResponse;
 import io.netty.handler.codec.http.LastHttpContent;
-import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
-import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
-import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
-import io.netty.handler.codec.http.websocketx.PongWebSocketFrame;
-import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.WebSocketFrame;
 
 import java.io.IOException;
@@ -123,15 +118,14 @@ public void handleRead(Channel channel, NettyResponseFuture<?> future, Object e)
         } else if (e instanceof WebSocketFrame) {
             final WebSocketFrame frame = (WebSocketFrame) e;
             WebSocketUpgradeHandler handler = (WebSocketUpgradeHandler) future.getAsyncHandler();
-            NettyWebSocket webSocket = (NettyWebSocket) handler.onCompleted();
+            NettyWebSocket webSocket = handler.onCompleted();
             // retain because we might buffer the frame
-            frame.retain();
-            if (handler.isOpen()) {
-                handleFrame(channel, frame, handler, webSocket);
+            if (webSocket.isReady()) {
+                webSocket.handleFrame(frame);
             } else {
                 // WebSocket hasn't been open yet, but upgrading the pipeline triggered a read and a frame was sent along the HTTP upgrade response
                 // as we want to keep sequential order (but can't notify user of open before upgrading so he doesn't to try send immediately), we have to buffer
-                handler.bufferFrame(() -> handleFrame(channel, frame, handler, webSocket));
+                webSocket.bufferFrame(frame);
             }
 
         } else if (!(e instanceof LastHttpContent)) {
@@ -140,36 +134,13 @@ public void handleRead(Channel channel, NettyResponseFuture<?> future, Object e)
         }
     }
 
-    private void handleFrame(Channel channel, WebSocketFrame frame, WebSocketUpgradeHandler handler, NettyWebSocket webSocket) {
-        if (frame instanceof TextWebSocketFrame) {
-            webSocket.onTextFrame((TextWebSocketFrame) frame);
-
-        } else if (frame instanceof BinaryWebSocketFrame) {
-            webSocket.onBinaryFrame((BinaryWebSocketFrame) frame);
-
-        } else if (frame instanceof CloseWebSocketFrame) {
-            Channels.setDiscard(channel);
-            CloseWebSocketFrame closeFrame = (CloseWebSocketFrame) frame;
-            webSocket.onClose(closeFrame.statusCode(), closeFrame.reasonText());
-            Channels.silentlyCloseChannel(channel);
-
-        } else if (frame instanceof PingWebSocketFrame) {
-            webSocket.onPing((PingWebSocketFrame) frame);
-
-        } else if (frame instanceof PongWebSocketFrame) {
-            webSocket.onPong((PongWebSocketFrame) frame);
-        }
-        // release because we had to retain in case the frame had to be buffered
-        frame.release();
-    }
-
     @Override
     public void handleException(NettyResponseFuture<?> future, Throwable e) {
         logger.warn("onError", e);
 
         try {
             WebSocketUpgradeHandler h = (WebSocketUpgradeHandler) future.getAsyncHandler();
-            NettyWebSocket webSocket = NettyWebSocket.class.cast(h.onCompleted());
+            NettyWebSocket webSocket = h.onCompleted();
             if (webSocket != null) {
                 webSocket.onError(e.getCause());
                 webSocket.close();
@@ -185,7 +156,7 @@ public void handleChannelInactive(NettyResponseFuture<?> future) {
 
         try {
             WebSocketUpgradeHandler h = (WebSocketUpgradeHandler) future.getAsyncHandler();
-            NettyWebSocket webSocket = NettyWebSocket.class.cast(h.onCompleted());
+            NettyWebSocket webSocket = h.onCompleted();
 
             logger.trace("Connection was closed abnormally (that is, with no close frame being received).");
             if (webSocket != null)
diff --git a/client/src/main/java/org/asynchttpclient/netty/ws/NettyWebSocket.java b/client/src/main/java/org/asynchttpclient/netty/ws/NettyWebSocket.java
@@ -23,12 +23,16 @@
 import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.PongWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
+import io.netty.handler.codec.http.websocketx.WebSocketFrame;
 
 import java.net.SocketAddress;
 import java.nio.charset.CharacterCodingException;
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.List;
 import java.util.concurrent.ConcurrentLinkedQueue;
 
+import org.asynchttpclient.netty.channel.Channels;
 import org.asynchttpclient.netty.util.Utf8ByteBufCharsetDecoder;
 import org.asynchttpclient.ws.WebSocket;
 import org.asynchttpclient.ws.WebSocketByteListener;
@@ -50,6 +54,9 @@ public class NettyWebSocket implements WebSocket {
     protected final Collection<WebSocketListener> listeners;
     private volatile boolean interestedInByteMessages;
     private volatile boolean interestedInTextMessages;
+    // no need for volatile because only mutated in IO thread
+    private boolean ready;
+    private List<WebSocketFrame> bufferedFrames;
 
     public NettyWebSocket(Channel channel, HttpHeaders upgradeHeaders) {
         this(channel, upgradeHeaders, new ConcurrentLinkedQueue<>());
@@ -61,6 +68,59 @@ public NettyWebSocket(Channel channel, HttpHeaders upgradeHeaders, Collection<We
         this.listeners = listeners;
     }
 
+    public boolean isReady() {
+        return ready;
+    }
+
+    public void bufferFrame(WebSocketFrame frame) {
+        if (bufferedFrames == null) {
+            bufferedFrames = new ArrayList<>(1);
+        }
+        frame.retain();
+        bufferedFrames.add(frame);
+    }
+
+    private void releaseBufferedFrames() {
+        for (WebSocketFrame frame : bufferedFrames) {
+            frame.release();
+        }
+    }
+
+    public void processBufferedFrames() {
+        ready = true;
+        if (bufferedFrames != null) {
+            try {
+                for (WebSocketFrame frame : bufferedFrames) {
+                    handleFrame(frame);
+                }
+            } finally {
+                releaseBufferedFrames();
+            }
+            bufferedFrames = null;
+        }
+    }
+
+    public void handleFrame(WebSocketFrame frame) {
+        if (frame instanceof TextWebSocketFrame) {
+            onTextFrame((TextWebSocketFrame) frame);
+
+        } else if (frame instanceof BinaryWebSocketFrame) {
+            onBinaryFrame((BinaryWebSocketFrame) frame);
+
+        } else if (frame instanceof CloseWebSocketFrame) {
+            Channels.setDiscard(channel);
+            CloseWebSocketFrame closeFrame = (CloseWebSocketFrame) frame;
+            onClose(closeFrame.statusCode(), closeFrame.reasonText());
+            Channels.silentlyCloseChannel(channel);
+
+        } else if (frame instanceof PingWebSocketFrame) {
+            onPing((PingWebSocketFrame) frame);
+
+        } else if (frame instanceof PongWebSocketFrame) {
+            onPong((PongWebSocketFrame) frame);
+        }
+    }
+
     @Override
     public HttpHeaders getUpgradeHeaders() {
         return upgradeHeaders;
@@ -189,6 +249,7 @@ public void close() {
     public void close(int statusCode, String reason) {
         onClose(statusCode, reason);
         listeners.clear();
+        releaseBufferedFrames();
     }
 
     public void onError(Throwable t) {
@@ -199,6 +260,7 @@ public void onError(Throwable t) {
                 LOGGER.error("WebSocketListener.onError crash", t2);
             }
         }
+        releaseBufferedFrames();
     }
 
     public void onClose(int code, String reason) {
diff --git a/client/src/main/java/org/asynchttpclient/ws/WebSocketUpgradeHandler.java b/client/src/main/java/org/asynchttpclient/ws/WebSocketUpgradeHandler.java
@@ -20,18 +20,17 @@
 import org.asynchttpclient.HttpResponseBodyPart;
 import org.asynchttpclient.HttpResponseHeaders;
 import org.asynchttpclient.HttpResponseStatus;
+import org.asynchttpclient.netty.ws.NettyWebSocket;
 
 /**
  * An {@link AsyncHandler} which is able to execute WebSocket upgrade. Use the Builder for configuring WebSocket options.
  */
-public class WebSocketUpgradeHandler implements AsyncHandler<WebSocket> {
+public class WebSocketUpgradeHandler implements AsyncHandler<NettyWebSocket> {
 
     private static final int SWITCHING_PROTOCOLS = io.netty.handler.codec.http.HttpResponseStatus.SWITCHING_PROTOCOLS.code();
 
-    private WebSocket webSocket;
-    private boolean open;
+    private NettyWebSocket webSocket;
     private final List<WebSocketListener> listeners;
-    private List<Runnable> bufferedFrames;
 
     public WebSocketUpgradeHandler(List<WebSocketListener> listeners) {
         this.listeners = listeners;
@@ -53,7 +52,7 @@ public final State onBodyPartReceived(HttpResponseBodyPart bodyPart) throws Exce
     }
 
     @Override
-    public final WebSocket onCompleted() throws Exception {
+    public final NettyWebSocket onCompleted() throws Exception {
         return webSocket;
     }
 
@@ -67,33 +66,16 @@ public final void onThrowable(Throwable t) {
         }
     }
 
-    public final void setWebSocket(WebSocket webSocket) {
+    public final void setWebSocket(NettyWebSocket webSocket) {
         this.webSocket = webSocket;
     }
     
     public final void onOpen() {
-        open = true;
         for (WebSocketListener listener : listeners) {
             webSocket.addWebSocketListener(listener);
             listener.onOpen(webSocket);
         }
-        if (bufferedFrames != null) {
-            for (Runnable bufferedFrame : bufferedFrames) {
-                bufferedFrame.run();
-            }
-            bufferedFrames = null;
-        }
-    }
-
-    public final boolean isOpen() {
-        return open;
-    }
-
-    public final void bufferFrame(Runnable bufferedFrame) {
-        if (bufferedFrames == null) {
-            bufferedFrames = new ArrayList<>(1);
-        }
-        bufferedFrames.add(bufferedFrame);
+        webSocket.processBufferedFrames();
     }
 
     /**
