428474 - Expose batch mode in the Jetty WebSocket API

5 files changed, 41 insertions, 8 deletions

diff --git a/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketRemoteEndpoint.java b/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketRemoteEndpoint.java
index 6d9b93dc82..9468ee6c57 100644
--- a/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketRemoteEndpoint.java
+++ b/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketRemoteEndpoint.java
@@ -436,9 +436,7 @@ public class WebSocketRemoteEndpoint implements RemoteEndpoint
         return batchMode;
     }
 
-    // Only the JSR needs to have this method exposed.
-    // In the Jetty implementation the batching is set
-    // at the moment of opening the session.
+    @Override
     public void setBatchMode(BatchMode batchMode)
     {
         this.batchMode = batchMode;
diff --git a/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketSession.java b/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketSession.java
index 3654c59c75..e0e2c91d08 100644
--- a/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketSession.java
+++ b/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketSession.java
@@ -71,6 +71,7 @@ public class WebSocketSession extends ContainerLifeCycle implements Session, Inc
     private WebSocketPolicy policy;
     private UpgradeRequest upgradeRequest;
     private UpgradeResponse upgradeResponse;
+    private BatchMode batchMode = BatchMode.AUTO;
 
     public WebSocketSession(URI requestURI, EventDriver websocket, LogicalConnection connection, SessionListener... sessionListeners)
     {
@@ -400,14 +401,20 @@ public class WebSocketSession extends ContainerLifeCycle implements Session, Inc
         }
         
         ClassLoader old = Thread.currentThread().getContextClassLoader();
-        try {
+        try 
+        {
             Thread.currentThread().setContextClassLoader(classLoader);
 
             // Upgrade success
             connection.getIOState().onConnected();
-    
+            
             // Connect remote
-            remote = new WebSocketRemoteEndpoint(connection,outgoingHandler,getBatchMode());
+            BatchMode endpointBatchMode = websocket.getBatchMode();
+            if (endpointBatchMode == null)
+            {
+                endpointBatchMode = this.getBatchMode();
+            }
+            remote = new WebSocketRemoteEndpoint(connection,outgoingHandler,endpointBatchMode);
 
             // Open WebSocket
             websocket.openSession(this);
@@ -496,11 +503,20 @@ public class WebSocketSession extends ContainerLifeCycle implements Session, Inc
     }
 
     /**
-     * @return the default (initial) value for the batching mode.
+     * @return the batching mode default for RemoteEndpoint behavior
      */
     public BatchMode getBatchMode()
     {
-        return BatchMode.AUTO;
+        return batchMode;
+    }
+    
+    /**
+     * Set the batch mode default for the RemoteEndpoint behavior.
+     * @param mode the batching mode.
+     */
+    public void setBatchMode(BatchMode mode)
+    {
+        this.batchMode = mode;
     }
 
     @Override
diff --git a/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/AbstractEventDriver.java b/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/AbstractEventDriver.java
index d446264743..df26e7e8f4 100644
--- a/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/AbstractEventDriver.java
+++ b/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/AbstractEventDriver.java
@@ -25,6 +25,7 @@ import org.eclipse.jetty.util.BufferUtil;
 import org.eclipse.jetty.util.Utf8Appendable.NotUtf8Exception;
 import org.eclipse.jetty.util.log.Log;
 import org.eclipse.jetty.util.log.Logger;
+import org.eclipse.jetty.websocket.api.BatchMode;
 import org.eclipse.jetty.websocket.api.CloseException;
 import org.eclipse.jetty.websocket.api.StatusCode;
 import org.eclipse.jetty.websocket.api.WebSocketPolicy;
@@ -208,6 +209,12 @@ public abstract class AbstractEventDriver implements IncomingFrames, EventDriver
     {
         /* TODO: provide annotation in future */
     }
+    
+    @Override
+    public BatchMode getBatchMode()
+    {
+        return null;
+    }
 
     @Override
     public void openSession(WebSocketSession session)
diff --git a/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/EventDriver.java b/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/EventDriver.java
index e4d4d836bf..3c3f28c5f1 100644
--- a/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/EventDriver.java
+++ b/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/EventDriver.java
@@ -23,6 +23,7 @@ import java.io.InputStream;
 import java.io.Reader;
 import java.nio.ByteBuffer;
 
+import org.eclipse.jetty.websocket.api.BatchMode;
 import org.eclipse.jetty.websocket.api.WebSocketPolicy;
 import org.eclipse.jetty.websocket.api.extensions.Frame;
 import org.eclipse.jetty.websocket.api.extensions.IncomingFrames;
@@ -34,6 +35,8 @@ public interface EventDriver extends IncomingFrames
     public WebSocketPolicy getPolicy();
 
     public WebSocketSession getSession();
+    
+    public BatchMode getBatchMode();
 
     public void onBinaryFrame(ByteBuffer buffer, boolean fin) throws IOException;
 
diff --git a/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/JettyAnnotatedEventDriver.java b/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/JettyAnnotatedEventDriver.java
index 45dcaf2b70..4d17bd0321 100644
--- a/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/JettyAnnotatedEventDriver.java
+++ b/jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/JettyAnnotatedEventDriver.java
@@ -23,6 +23,7 @@ import java.io.InputStream;
 import java.io.Reader;
 import java.nio.ByteBuffer;
 
+import org.eclipse.jetty.websocket.api.BatchMode;
 import org.eclipse.jetty.websocket.api.WebSocketPolicy;
 import org.eclipse.jetty.websocket.api.annotations.WebSocket;
 import org.eclipse.jetty.websocket.api.extensions.Frame;
@@ -40,6 +41,7 @@ public class JettyAnnotatedEventDriver extends AbstractEventDriver
 {
     private final JettyAnnotatedMetadata events;
     private boolean hasCloseBeenCalled = false;
+    private BatchMode batchMode;
 
     public JettyAnnotatedEventDriver(WebSocketPolicy policy, Object websocket, JettyAnnotatedMetadata events)
     {
@@ -64,6 +66,13 @@ public class JettyAnnotatedEventDriver extends AbstractEventDriver
         {
             this.policy.setIdleTimeout(anno.maxIdleTime());
         }
+        this.batchMode = anno.batchMode();
+    }
+    
+    @Override
+    public BatchMode getBatchMode()
+    {
+        return this.batchMode;
     }
 
     @Override