diff --git a/src/main/java/io/vertx/core/http/impl/WebSocketImplBase.java b/src/main/java/io/vertx/core/http/impl/WebSocketImplBase.java index 7467dadaa2f..aca102c6d99 100644 --- a/src/main/java/io/vertx/core/http/impl/WebSocketImplBase.java +++ b/src/main/java/io/vertx/core/http/impl/WebSocketImplBase.java @@ -75,6 +75,7 @@ public abstract class WebSocketImplBase implements WebS private Object metric; private Handler handler; private Handler frameHandler; + private FrameAggregator frameAggregator; private Handler pongHandler; private Handler drainHandler; private Handler exceptionHandler; @@ -525,9 +526,14 @@ protected void handleClose(boolean graceful) { } private void receiveFrame(WebSocketFrameInternal frame) { + Handler frameAggregator; Handler frameHandler; synchronized (conn) { frameHandler = this.frameHandler; + frameAggregator = this.frameAggregator; + } + if (frameAggregator != null) { + context.dispatch(frame, frameAggregator); } if (frameHandler != null) { context.dispatch(frame, frameHandler); @@ -670,10 +676,21 @@ public S frameHandler(Handler handler) { public WebSocketBase textMessageHandler(Handler handler) { synchronized (conn) { checkClosed(); - if (frameHandler == null || frameHandler.getClass() != FrameAggregator.class) { - frameHandler = new FrameAggregator(); + if (handler != null) { + if (frameAggregator == null) { + frameAggregator = new FrameAggregator(); + } + frameAggregator.textMessageHandler = handler; + } else { + if (frameAggregator != null) { + if (frameAggregator.binaryMessageHandler == null) { + frameAggregator = null; + } else { + frameAggregator.textMessageHandler = null; + frameAggregator.textMessageBuffer = null; + } + } } - ((FrameAggregator) frameHandler).textMessageHandler = handler; return this; } } @@ -682,10 +699,21 @@ public WebSocketBase textMessageHandler(Handler handler) { public S binaryMessageHandler(Handler handler) { synchronized (conn) { checkClosed(); - if (frameHandler == null || frameHandler.getClass() != FrameAggregator.class) { - frameHandler = new FrameAggregator(); + if (handler != null) { + if (frameAggregator == null) { + frameAggregator = new FrameAggregator(); + } + frameAggregator.binaryMessageHandler = handler; + } else { + if (frameAggregator != null) { + if (frameAggregator.textMessageHandler == null) { + frameAggregator = null; + } else { + frameAggregator.binaryMessageHandler = null; + frameAggregator.binaryMessageBuffer = null; + } + } } - ((FrameAggregator) frameHandler).binaryMessageHandler = handler; return (S) this; } }