Issue #4502 - onClose can now be triggered on receiving a close frame

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>
jetty · Jan 29, 2020 · 4ef208e · 4ef208e
1 parent f6fd3c4
commit 4ef208e
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 14 deletions.
diff --git a/...on/src/main/java/org/eclipse/jetty/websocket/javax/common/JavaxWebSocketFrameHandler.java b/...on/src/main/java/org/eclipse/jetty/websocket/javax/common/JavaxWebSocketFrameHandler.java
@@ -22,10 +22,12 @@
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
 import java.nio.ByteBuffer;
+import java.nio.channels.ClosedChannelException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 import javax.websocket.ClientEndpointConfig;
 import javax.websocket.CloseReason;
@@ -60,6 +62,8 @@ public class JavaxWebSocketFrameHandler implements FrameHandler
     private final Logger logger;
     private final JavaxWebSocketContainer container;
     private final Object endpointInstance;
+    private final AtomicBoolean closeNotified = new AtomicBoolean();
+
     /**
      * List of configured named variables in the uri-template.
      * <p>
@@ -278,24 +282,41 @@ public void onFrame(Frame frame, Callback callback)
             dataType = OpCode.UNDEFINED;
     }
 
+    public void onClose(Frame frame, Callback callback)
+    {
+        notifyOnClose(CloseStatus.getCloseStatus(frame), callback);
+    }
+
     @Override
     public void onClosed(CloseStatus closeStatus, Callback callback)
     {
+        notifyOnClose(closeStatus, callback);
+        container.notifySessionListeners((listener) -> listener.onJavaxWebSocketSessionClosed(session));
+    }
+
+    private void notifyOnClose(CloseStatus closeStatus, Callback callback)
+    {
+        // Make sure onClose is only notified once.
+        if (!closeNotified.compareAndSet(false, true))
+        {
+            callback.failed(new ClosedChannelException());
+            return;
+        }
+
         try
         {
             if (closeHandle != null)
             {
                 CloseReason closeReason = new CloseReason(CloseReason.CloseCodes.getCloseCode(closeStatus.getCode()), closeStatus.getReason());
                 closeHandle.invoke(closeReason);
             }
+
             callback.succeeded();
         }
         catch (Throwable cause)
         {
             callback.failed(new WebSocketException(endpointInstance.getClass().getSimpleName() + " CLOSE method error: " + cause.getMessage(), cause));
         }
-
-        container.notifySessionListeners((listener) -> listener.onJavaxWebSocketSessionClosed(session));
     }
 
     @Override
@@ -572,11 +593,6 @@ private void acceptMessage(Frame frame, Callback callback)
             activeMessageSink = null;
     }
 
-    public void onClose(Frame frame, Callback callback)
-    {
-        callback.succeeded();
-    }
-
     public void onPing(Frame frame, Callback callback)
     {
         ByteBuffer payload = BufferUtil.copy(frame.getPayload());

diff --git a/...y-common/src/main/java/org/eclipse/jetty/websocket/common/JettyWebSocketFrameHandler.java b/...y-common/src/main/java/org/eclipse/jetty/websocket/common/JettyWebSocketFrameHandler.java
@@ -20,7 +20,9 @@
 
 import java.lang.invoke.MethodHandle;
 import java.nio.ByteBuffer;
+import java.nio.channels.ClosedChannelException;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.eclipse.jetty.util.BufferUtil;
 import org.eclipse.jetty.util.Callback;
@@ -57,6 +59,7 @@ private enum SuspendState
     private final WebSocketContainer container;
     private final Object endpointInstance;
     private final BatchMode batchMode;
+    private final AtomicBoolean closeNotified = new AtomicBoolean();
     private MethodHandle openHandle;
     private MethodHandle closeHandle;
     private MethodHandle errorHandle;
@@ -278,6 +281,11 @@ public void onError(Throwable cause, Callback callback)
         }
     }
 
+    private void onCloseFrame(Frame frame, Callback callback)
+    {
+        notifyOnClose(CloseStatus.getCloseStatus(frame), callback);
+    }
+
     @Override
     public void onClosed(CloseStatus closeStatus, Callback callback)
     {
@@ -287,6 +295,19 @@ public void onClosed(CloseStatus closeStatus, Callback callback)
             state = SuspendState.CLOSED;
         }
 
+        notifyOnClose(closeStatus, callback);
+        container.notifySessionListeners((listener) -> listener.onWebSocketSessionClosed(session));
+    }
+
+    private void notifyOnClose(CloseStatus closeStatus, Callback callback)
+    {
+        // Make sure onClose is only notified once.
+        if (!closeNotified.compareAndSet(false, true))
+        {
+            callback.failed(new ClosedChannelException());
+            return;
+        }
+
         try
         {
             if (closeHandle != null)
@@ -298,8 +319,6 @@ public void onClosed(CloseStatus closeStatus, Callback callback)
         {
             callback.failed(new WebSocketException(endpointInstance.getClass().getSimpleName() + " CLOSE method error: " + cause.getMessage(), cause));
         }
-
-        container.notifySessionListeners((listener) -> listener.onWebSocketSessionClosed(session));
     }
 
     public String toString()
@@ -330,11 +349,6 @@ private void onBinaryFrame(Frame frame, Callback callback)
         acceptMessage(frame, callback);
     }
 
-    private void onCloseFrame(Frame frame, Callback callback)
-    {
-        callback.succeeded();
-    }
-
     private void onContinuationFrame(Frame frame, Callback callback)
     {
         acceptMessage(frame, callback);