Issue #207 - Careful distinction between container policy, vs session policy

# Conflicts:
#	jetty-websocket/websocket-client/src/main/java/org/eclipse/jetty/websocket/client/io/UpgradeConnection.java

jetty-websocket/javax-websocket-client-impl/src/main/java/org/eclipse/jetty/websocket/jsr356/ClientContainer.java

@@ -103,13 +103,14 @@ public class ClientContainer extends ContainerLifeCycle implements WebSocketCont
     }
     
     public EndpointFunctions newJsrEndpointFunction(Object endpoint,
+                                                    WebSocketPolicy sessionPolicy,
                                                     AvailableEncoders availableEncoders,
                                                     AvailableDecoders availableDecoders,
                                                     Map<String, String> pathParameters,
                                                     EndpointConfig config)
     {
         return new JsrEndpointFunctions(endpoint,
-                getPolicy(),
+                sessionPolicy,
                 getExecutor(),
                 availableEncoders,
                 availableDecoders,

jetty-websocket/javax-websocket-client-impl/src/main/java/org/eclipse/jetty/websocket/jsr356/JsrSession.java

@@ -42,7 +42,6 @@ import javax.websocket.WebSocketContainer;
 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.WebSocketPolicy;
 import org.eclipse.jetty.websocket.api.extensions.ExtensionConfig;
 import org.eclipse.jetty.websocket.common.LogicalConnection;
 import org.eclipse.jetty.websocket.common.WebSocketSession;
@@ -71,9 +70,9 @@ public class JsrSession extends WebSocketSession implements javax.websocket.Sess
     private JsrAsyncRemote asyncRemote;
     private JsrBasicRemote basicRemote;
     
-    public JsrSession(ClientContainer container, String id, URI requestURI, Object websocket, WebSocketPolicy policy, LogicalConnection connection)
+    public JsrSession(ClientContainer container, String id, URI requestURI, Object websocket, LogicalConnection connection)
     {
-        super(container, requestURI, websocket, policy, connection);
+        super(container, requestURI, websocket, connection);
         
         this.container = container;
         
@@ -100,6 +99,7 @@ public class JsrSession extends WebSocketSession implements javax.websocket.Sess
         // Delegate to container to obtain correct version of JsrEndpointFunctions
         // Could be a Client version, or a Server version
         return container.newJsrEndpointFunction(endpoint,
+                getPolicy(),
                 availableEncoders,
                 availableDecoders,
                 pathParameters,

jetty-websocket/javax-websocket-client-impl/src/main/java/org/eclipse/jetty/websocket/jsr356/JsrSessionFactory.java

@@ -20,7 +20,6 @@ package org.eclipse.jetty.websocket.jsr356;
 
 import java.net.URI;
 
-import org.eclipse.jetty.websocket.api.WebSocketPolicy;
 import org.eclipse.jetty.websocket.common.LogicalConnection;
 import org.eclipse.jetty.websocket.common.SessionFactory;
 import org.eclipse.jetty.websocket.common.WebSocketSession;
@@ -35,9 +34,9 @@ public class JsrSessionFactory implements SessionFactory
     }
 
     @Override
-    public WebSocketSession createSession(URI requestURI, Object websocket, WebSocketPolicy policy, LogicalConnection connection)
+    public WebSocketSession createSession(URI requestURI, Object websocket, LogicalConnection connection)
     {
-        return new JsrSession(container,connection.getId(),requestURI,websocket,policy,connection);
+        return new JsrSession(container,connection.getId(),requestURI,websocket,connection);
     }
 
     @Override

jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/JsrSessionTest.java

@@ -63,7 +63,7 @@ public class JsrSessionTest
         ConfiguredEndpoint ei = new ConfiguredEndpoint(new DummyEndpoint(), config);
         
         // Session
-        session = new JsrSession(container, id, requestURI, ei, policy, connection);
+        session = new JsrSession(container, id, requestURI, ei, connection);
         session.start();
     }
     

jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/function/JsrEndpointFunctions_OnCloseTest.java

@@ -75,7 +75,7 @@ public class JsrEndpointFunctions_OnCloseTest
         DummyConnection connection = new DummyConnection(policy);
         ClientEndpointConfig config = new EmptyClientEndpointConfig();
         ConfiguredEndpoint ei = new ConfiguredEndpoint(websocket, config);
-        return new JsrSession(container, id, requestURI, ei, policy, connection);
+        return new JsrSession(container, id, requestURI, ei, connection);
     }
 
     private void assertOnCloseInvocation(TrackingSocket socket, String expectedEventFormat, Object... args) throws Exception

jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/function/JsrEndpointFunctions_OnErrorTest.java

@@ -72,7 +72,7 @@ public class JsrEndpointFunctions_OnErrorTest
         DummyConnection connection = new DummyConnection(policy);
         ClientEndpointConfig config = new EmptyClientEndpointConfig();
         ConfiguredEndpoint ei = new ConfiguredEndpoint(websocket, config);
-        return new JsrSession(container, id, requestURI, ei, policy, connection);
+        return new JsrSession(container, id, requestURI, ei, connection);
     }
 
     private void assertOnErrorInvocation(TrackingSocket socket, String expectedEventFormat, Object... args) throws Exception

jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/function/JsrEndpointFunctions_OnMessage_BinaryStreamTest.java

@@ -82,7 +82,7 @@ public class JsrEndpointFunctions_OnMessage_BinaryStreamTest
         DummyConnection connection = new DummyConnection(policy);
         ClientEndpointConfig config = new EmptyClientEndpointConfig();
         ConfiguredEndpoint ei = new ConfiguredEndpoint(websocket, config);
-        return new JsrSession(container, id, requestURI, ei, policy, connection);
+        return new JsrSession(container, id, requestURI, ei, connection);
     }
     
     @SuppressWarnings("Duplicates")

jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/function/JsrEndpointFunctions_OnMessage_BinaryTest.java

@@ -84,7 +84,7 @@ public class JsrEndpointFunctions_OnMessage_BinaryTest
         DummyConnection connection = new DummyConnection(policy);
         ClientEndpointConfig config = new EmptyClientEndpointConfig();
         ConfiguredEndpoint ei = new ConfiguredEndpoint(websocket, config);
-        return new JsrSession(container, id, requestURI, ei, policy, connection);
+        return new JsrSession(container, id, requestURI, ei, connection);
     }
     
     private void assertOnMessageInvocation(TrackingSocket socket, String expectedEventFormat, Object... args) throws Exception

jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/function/JsrEndpointFunctions_OnMessage_TextStreamTest.java

@@ -82,7 +82,7 @@ public class JsrEndpointFunctions_OnMessage_TextStreamTest
         DummyConnection connection = new DummyConnection(policy);
         ClientEndpointConfig config = new EmptyClientEndpointConfig();
         ConfiguredEndpoint ei = new ConfiguredEndpoint(websocket, config);
-        return new JsrSession(container, id, requestURI, ei, policy, connection);
+        return new JsrSession(container, id, requestURI, ei, connection);
     }
     
     @SuppressWarnings("Duplicates")

jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/function/JsrEndpointFunctions_OnMessage_TextTest.java

@@ -82,7 +82,7 @@ public class JsrEndpointFunctions_OnMessage_TextTest
         DummyConnection connection = new DummyConnection(policy);
         ClientEndpointConfig config = new EmptyClientEndpointConfig();
         ConfiguredEndpoint ei = new ConfiguredEndpoint(websocket, config);
-        return new JsrSession(container, id, requestURI, ei, policy, connection);
+        return new JsrSession(container, id, requestURI, ei, connection);
     }
     
     private void onText(TrackingSocket socket, String msg) throws Exception

jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/function/JsrEndpointFunctions_OnOpenTest.java

@@ -71,7 +71,7 @@ public class JsrEndpointFunctions_OnOpenTest
         DummyConnection connection = new DummyConnection(policy);
         ClientEndpointConfig config = new EmptyClientEndpointConfig();
         ConfiguredEndpoint ei = new ConfiguredEndpoint(websocket, config);
-        return new JsrSession(container, id, requestURI, ei, policy, connection);
+        return new JsrSession(container, id, requestURI, ei, connection);
     }
 
     private void assertOnOpenInvocation(TrackingSocket socket, String expectedEventFormat, Object... args) throws Exception

jetty-websocket/javax-websocket-server-impl/src/main/java/org/eclipse/jetty/websocket/jsr356/server/ServerContainer.java

@@ -39,6 +39,7 @@ import org.eclipse.jetty.servlet.ServletContextHandler;
 import org.eclipse.jetty.util.annotation.ManagedObject;
 import org.eclipse.jetty.util.log.Log;
 import org.eclipse.jetty.util.log.Logger;
+import org.eclipse.jetty.websocket.api.WebSocketPolicy;
 import org.eclipse.jetty.websocket.common.WebSocketSession;
 import org.eclipse.jetty.websocket.common.function.EndpointFunctions;
 import org.eclipse.jetty.websocket.jsr356.ClientContainer;
@@ -187,7 +188,7 @@ public class ServerContainer extends ClientContainer implements javax.websocket.
                 pathParameters.put(variable, "0");
             }
             
-            endpointFunctions = newJsrEndpointFunction(endpoint, availableEncoders, availableDecoders, pathParameters, config);
+            endpointFunctions = newJsrEndpointFunction(endpoint, getPolicy(), availableEncoders, availableDecoders, pathParameters, config);
             endpointFunctions.start(); // this should trigger an exception if endpoint is invalid.
         }
         catch (InstantiationException e)
@@ -221,13 +222,14 @@ public class ServerContainer extends ClientContainer implements javax.websocket.
     
     @Override
     public EndpointFunctions newJsrEndpointFunction(Object endpoint,
+                                                    WebSocketPolicy sessionPolicy,
                                                     AvailableEncoders availableEncoders,
                                                     AvailableDecoders availableDecoders,
                                                     Map<String, String> pathParameters,
                                                     EndpointConfig config)
     {
         return new JsrServerEndpointFunctions(endpoint,
-                getPolicy(),
+                sessionPolicy,
                 getExecutor(),
                 availableEncoders,
                 availableDecoders,

jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/AnnotatedServerEndpointTest.java

@@ -22,7 +22,7 @@ import static org.hamcrest.Matchers.containsString;
 
 import java.net.URI;
 import java.nio.file.Path;
-import java.util.Queue;
+import java.util.List;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
@@ -82,6 +82,7 @@ public class AnnotatedServerEndpointTest
         {
             client.start();
             JettyEchoSocket clientEcho = new JettyEchoSocket();
+            Future<List<String>> clientMessagesFuture = clientEcho.expectedMessages(1);
             URI uri = server.getServerBaseURI().resolve("echo");
             ClientUpgradeRequest req = new ClientUpgradeRequest();
             req.setSubProtocols("echo");
@@ -90,9 +91,9 @@ public class AnnotatedServerEndpointTest
             foo.get(1,TimeUnit.SECONDS);
 
             clientEcho.sendMessage(message);
-            Queue<String> msgs = clientEcho.awaitMessages(1);
+            List<String> msgs = clientMessagesFuture.get(1, TimeUnit.SECONDS);
 
-            String response = msgs.poll();
+            String response = msgs.get(0);
             for (String expected : expectedTexts)
             {
                 Assert.assertThat("Expected message",response,containsString(expected));

jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/BasicEndpointTest.java

@@ -19,7 +19,7 @@
 package org.eclipse.jetty.websocket.jsr356.server;
 
 import java.net.URI;
-import java.util.Queue;
+import java.util.List;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
@@ -69,12 +69,13 @@ public class BasicEndpointTest
             {
                 client.start();
                 JettyEchoSocket clientEcho = new JettyEchoSocket();
+                Future<List<String>> clientMessagesFuture = clientEcho.expectedMessages(1);
                 Future<Session> future = client.connect(clientEcho,uri.resolve("echo"));
                 // wait for connect
                 future.get(1,TimeUnit.SECONDS);
                 clientEcho.sendMessage("Hello World");
-                Queue<String> msgs = clientEcho.awaitMessages(1);
+                List<String> msgs = clientMessagesFuture.get(1, TimeUnit.SECONDS);
-                Assert.assertEquals("Expected message","Hello World",msgs.poll());
+                Assert.assertEquals("Expected message","Hello World",msgs.get(0));
             }
             finally
             {

jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/IdleTimeoutTest.java

@@ -18,12 +18,14 @@
 
 package org.eclipse.jetty.websocket.jsr356.server;
 
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertThat;
 
 import java.net.URI;
-import java.util.Queue;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -32,7 +34,9 @@ import org.eclipse.jetty.toolchain.test.MavenTestingUtils;
 import org.eclipse.jetty.util.log.Log;
 import org.eclipse.jetty.util.log.Logger;
 import org.eclipse.jetty.webapp.WebAppContext;
+import org.eclipse.jetty.websocket.api.CloseException;
 import org.eclipse.jetty.websocket.api.Session;
+import org.eclipse.jetty.websocket.api.StatusCode;
 import org.eclipse.jetty.websocket.client.WebSocketClient;
 import org.eclipse.jetty.websocket.common.test.LeakTrackingBufferPoolRule;
 import org.eclipse.jetty.websocket.jsr356.server.samples.idletimeout.IdleTimeoutContextListener;
@@ -80,6 +84,8 @@ public class IdleTimeoutTest
         {
             client.start();
             JettyEchoSocket clientEcho = new JettyEchoSocket();
+            Future<List<String>> clientMessagesFuture = clientEcho.expectedMessages(1);
+            
             if (LOG.isDebugEnabled())
                 LOG.debug("Client Attempting to connect");
             Future<Session> future = client.connect(clientEcho,uri);
@@ -93,18 +99,26 @@ public class IdleTimeoutTest
             TimeUnit.SECONDS.sleep(1);
             if (LOG.isDebugEnabled())
                 LOG.debug("Waited 1 second");
-            if (clientEcho.getClosed() == false)
+            
+            // Try to write
+            clientEcho.sendMessage("You shouldn't be there");
+            try
             {
-                // Try to write
+                List<String> msgs = clientMessagesFuture.get(1, TimeUnit.SECONDS);
-                clientEcho.sendMessage("You shouldn't be there");
+                assertThat("Should not have received messages echoed back",msgs,is(empty()));
-                try
+            }
+            catch (ExecutionException e)
+            {
+                Throwable cause = e.getCause();
+                if(cause instanceof CloseException)
                 {
-                    Queue<String> msgs = clientEcho.awaitMessages(1);
+                    CloseException ce = (CloseException) cause;
-                    assertThat("Should not have received messages echoed back",msgs,is(empty()));
+                    assertThat("CloseException.statusCode", ce.getStatusCode(), is(StatusCode.SHUTDOWN));
+                    assertThat("CloseException.reason", ce.getMessage(), containsString("Idle Timeout"));
                 }
-                catch (TimeoutException | InterruptedException e)
+                else
                 {
-                    // valid success path
+                    throw e;
                 }
             }
         }

jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/JettyEchoSocket.java

@@ -19,18 +19,19 @@
 package org.eclipse.jetty.websocket.jsr356.server;
 
 import java.io.IOException;
-import java.util.Queue;
+import java.util.ArrayList;
-import java.util.concurrent.TimeUnit;
+import java.util.List;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.locks.Lock;
+import java.util.concurrent.Future;
-import java.util.concurrent.locks.ReentrantLock;
+import java.util.concurrent.atomic.AtomicInteger;
 
-import org.eclipse.jetty.toolchain.test.EventQueue;
 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.RemoteEndpoint;
 import org.eclipse.jetty.websocket.api.Session;
+import org.eclipse.jetty.websocket.api.StatusCode;
 import org.eclipse.jetty.websocket.api.annotations.OnWebSocketClose;
 import org.eclipse.jetty.websocket.api.annotations.OnWebSocketConnect;
 import org.eclipse.jetty.websocket.api.annotations.OnWebSocketError;
@@ -46,92 +47,73 @@ import org.eclipse.jetty.websocket.api.annotations.WebSocket;
 public class JettyEchoSocket
 {
     private static final Logger LOG = Log.getLogger(JettyEchoSocket.class);
-    @SuppressWarnings("unused")
-    private Session session;
-    private Lock remoteLock = new ReentrantLock();
     private RemoteEndpoint remote;
-    private EventQueue<String> incomingMessages = new EventQueue<>();
+    private CompletableFuture<List<String>> expectedMessagesFuture;
-
+    private AtomicInteger expectedMessageCount;
-    public Queue<String> awaitMessages(int expected) throws TimeoutException, InterruptedException
+    private List<String> messages = new ArrayList<>();
+    
+    public Future<List<String>> expectedMessages(int expected)
     {
-        incomingMessages.awaitEventCount(expected,2,TimeUnit.SECONDS);
+        expectedMessagesFuture = new CompletableFuture<>();
-        return incomingMessages;
+        expectedMessageCount = new AtomicInteger(expected);
+        return expectedMessagesFuture;
     }
-
+    
-    public boolean getClosed()
-    {
-        remoteLock.lock();
-        try
-        {
-            return (remote == null);
-        }
-        finally
-        {
-            remoteLock.unlock();
-        }
-    }
-
     @OnWebSocketClose
     public void onClose(int code, String reason)
     {
-        session = null;
+        remote = null;
-        remoteLock.lock();
+        synchronized (expectedMessagesFuture)
-        try
         {
-            remote = null;
+            if ((code != StatusCode.NORMAL) ||
-        }
+                    (code != StatusCode.NO_CODE))
-        finally
+            {
-        {
+                expectedMessagesFuture.completeExceptionally(new CloseException(code, reason));
-            remoteLock.unlock();
+            }
         }
     }
-
+    
     @OnWebSocketError
     public void onError(Throwable t)
     {
         LOG.warn(t);
+        synchronized (expectedMessagesFuture)
+        {
+            expectedMessagesFuture.completeExceptionally(t);
+        }
     }
-
+    
     @OnWebSocketMessage
     public void onMessage(String msg) throws IOException
     {
-        incomingMessages.add(msg);
+        messages.add(msg);
+        synchronized (expectedMessagesFuture)
+        {
+            int countLeft = expectedMessageCount.decrementAndGet();
+            if (countLeft <= 0)
+            {
+                expectedMessagesFuture.complete(messages);
+            }
+        }
         sendMessage(msg);
     }
-
+    
     @OnWebSocketConnect
     public void onOpen(Session session)
     {
-        this.session = session;
+        this.remote = session.getRemote();
-        remoteLock.lock();
-        try
-        {
-            this.remote = session.getRemote();
-        }
-        finally
-        {
-            remoteLock.unlock();
-        }
     }
-
+    
     public void sendMessage(String msg) throws IOException
     {
-        remoteLock.lock();
+        RemoteEndpoint r = remote;
-        try
+        if (r == null)
         {
-            RemoteEndpoint r = remote;
+            return;
-            if (r == null)
-            {
-                return;
-            }
-
-            r.sendStringByFuture(msg);
-            if (r.getBatchMode() == BatchMode.ON)
-                r.flush();
-        }
-        finally
-        {
-            remoteLock.unlock();
         }
+        
+        r.sendStringByFuture(msg);
+        if (r.getBatchMode() == BatchMode.ON)
+            r.flush();
     }
 }

jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/JsrServerEndpointFunctions_OnMessage_TextStreamTest.java

@@ -83,7 +83,7 @@ public class JsrServerEndpointFunctions_OnMessage_TextStreamTest
         DummyConnection connection = new DummyConnection(policy);
         ClientEndpointConfig config = new EmptyClientEndpointConfig();
         ConfiguredEndpoint ei = new ConfiguredEndpoint(websocket, config);
-        return new JsrSession(container, id, requestURI, ei, policy, connection);
+        return new JsrSession(container, id, requestURI, ei, connection);
     }
     
     @SuppressWarnings("Duplicates")

jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/LargeAnnotatedTest.java

@@ -21,7 +21,7 @@ package org.eclipse.jetty.websocket.jsr356.server;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
-import java.util.Queue;
+import java.util.List;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
@@ -48,6 +48,7 @@ public class LargeAnnotatedTest
     @Rule
     public LeakTrackingBufferPoolRule bufferPool = new LeakTrackingBufferPoolRule("Test");
 
+    @SuppressWarnings("Duplicates")
     @Test
     public void testEcho() throws Exception
     {
@@ -70,6 +71,7 @@ public class LargeAnnotatedTest
                 client.getPolicy().setMaxTextMessageSize(128*1024);
                 client.start();
                 JettyEchoSocket clientEcho = new JettyEchoSocket();
+                Future<List<String>> clientMessagesFuture = clientEcho.expectedMessages(1);
                 Future<Session> foo = client.connect(clientEcho,uri.resolve("echo/large"));
                 // wait for connect
                 foo.get(1,TimeUnit.SECONDS);
@@ -78,8 +80,8 @@ public class LargeAnnotatedTest
                 Arrays.fill(txt,(byte)'o');
                 String msg = new String(txt,StandardCharsets.UTF_8);
                 clientEcho.sendMessage(msg);
-                Queue<String> msgs = clientEcho.awaitMessages(1);
+                List<String> msgs = clientMessagesFuture.get(1, TimeUnit.SECONDS);
-                Assert.assertEquals("Expected message",msg,msgs.poll());
+                Assert.assertEquals("Expected message",msg,msgs.get(0));
             }
             finally
             {

jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/LargeContainerTest.java

@@ -21,7 +21,7 @@ package org.eclipse.jetty.websocket.jsr356.server;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
-import java.util.Queue;
+import java.util.List;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
@@ -48,6 +48,7 @@ public class LargeContainerTest
     @Rule
     public LeakTrackingBufferPoolRule bufferPool = new LeakTrackingBufferPoolRule("Test");
 
+    @SuppressWarnings("Duplicates")
     @Test
     public void testEcho() throws Exception
     {
@@ -70,6 +71,7 @@ public class LargeContainerTest
                 client.getPolicy().setMaxTextMessageSize(128*1024);
                 client.start();
                 JettyEchoSocket clientEcho = new JettyEchoSocket();
+                Future<List<String>> clientMessagesFuture = clientEcho.expectedMessages(1);
                 Future<Session> foo = client.connect(clientEcho,uri.resolve("echo/large"));
                 // wait for connect
                 foo.get(1,TimeUnit.SECONDS);
@@ -78,8 +80,8 @@ public class LargeContainerTest
                 Arrays.fill(txt,(byte)'o');
                 String msg = new String(txt,StandardCharsets.UTF_8);
                 clientEcho.sendMessage(msg);
-                Queue<String> msgs = clientEcho.awaitMessages(1);
+                List<String> msgs = clientMessagesFuture.get(1, TimeUnit.SECONDS);
-                Assert.assertEquals("Expected message",msg,msgs.poll());
+                Assert.assertEquals("Expected message",msg,msgs.get(0));
             }
             finally
             {

jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/OnMessageReturnTest.java

@@ -19,7 +19,7 @@
 package org.eclipse.jetty.websocket.jsr356.server;
 
 import java.net.URI;
-import java.util.Queue;
+import java.util.List;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
@@ -61,12 +61,13 @@ public class OnMessageReturnTest
             {
                 client.start();
                 JettyEchoSocket clientEcho = new JettyEchoSocket();
+                Future<List<String>> clientMessagesFuture = clientEcho.expectedMessages(1);
                 Future<Session> future = client.connect(clientEcho,uri.resolve("echoreturn"));
                 // wait for connect
                 future.get(1,TimeUnit.SECONDS);
                 clientEcho.sendMessage("Hello World");
-                Queue<String> msgs = clientEcho.awaitMessages(1);
+                List<String> msgs = clientMessagesFuture.get(1, TimeUnit.SECONDS);
-                Assert.assertEquals("Expected message","Hello World",msgs.poll());
+                Assert.assertEquals("Expected message","Hello World",msgs.get(0));
             }
             finally
             {

jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/SessionTest.java

@@ -24,7 +24,6 @@ import java.net.URI;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
-import java.util.Queue;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -139,12 +138,13 @@ public class SessionTest
         {
             client.start();
             JettyEchoSocket clientEcho = new JettyEchoSocket();
+            Future<List<String>> clientMessagesFuture = clientEcho.expectedMessages(1);
             Future<Session> future = client.connect(clientEcho,serverUri.resolve(requestPath));
             // wait for connect
             future.get(1,TimeUnit.SECONDS);
             clientEcho.sendMessage(requestMessage);
-            Queue<String> msgs = clientEcho.awaitMessages(1);
+            List<String> msgs = clientMessagesFuture.get(1, TimeUnit.SECONDS);
-            Assert.assertThat("Expected message",msgs.poll(),is(expectedResponse));
+            Assert.assertThat("Expected message",msgs.get(0),is(expectedResponse));
         }
         finally
         {

jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/LogicalConnection.java

@@ -23,6 +23,7 @@ import java.util.concurrent.Executor;
 
 import org.eclipse.jetty.io.ByteBufferPool;
 import org.eclipse.jetty.websocket.api.SuspendToken;
+import org.eclipse.jetty.websocket.api.WebSocketPolicy;
 import org.eclipse.jetty.websocket.api.extensions.IncomingFrames;
 import org.eclipse.jetty.websocket.api.extensions.OutgoingFrames;
 import org.eclipse.jetty.websocket.common.io.IOState;
@@ -97,6 +98,13 @@ public interface LogicalConnection extends OutgoingFrames, SuspendToken
      * @return the idle timeout in milliseconds
      */
     long getMaxIdleTimeout();
+    
+    /**
+     * Get the Connection based WebSocket Policy.
+     *
+     * @return the WebSocket policy for the connection
+     */
+    WebSocketPolicy getPolicy();
 
     /**
      * Get the remote Address in use for this connection.

jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/SessionFactory.java

@@ -20,8 +20,6 @@ package org.eclipse.jetty.websocket.common;
 
 import java.net.URI;
 
-import org.eclipse.jetty.websocket.api.WebSocketPolicy;
-
 /**
  * Interface for creating jetty {@link WebSocketSession} objects.
  */
@@ -39,9 +37,8 @@ public interface SessionFactory
      *
      * @param requestURI
      * @param websocket
-     * @param policy
      * @param connection
      * @return
      */
-    WebSocketSession createSession(URI requestURI, Object websocket, WebSocketPolicy policy, LogicalConnection connection);
+    WebSocketSession createSession(URI requestURI, Object websocket, LogicalConnection connection);
 }

jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketSession.java

@@ -100,7 +100,7 @@ public class WebSocketSession extends ContainerLifeCycle implements Session, Rem
     private UpgradeResponse upgradeResponse;
     private CompletableFuture<Session> openFuture;
 
-    public WebSocketSession(WebSocketContainerScope containerScope, URI requestURI, Object endpoint, WebSocketPolicy policy, LogicalConnection connection)
+    public WebSocketSession(WebSocketContainerScope containerScope, URI requestURI, Object endpoint, LogicalConnection connection)
     {
         Objects.requireNonNull(containerScope, "Container Scope cannot be null");
         Objects.requireNonNull(requestURI, "Request URI cannot be null");
@@ -113,7 +113,7 @@ public class WebSocketSession extends ContainerLifeCycle implements Session, Rem
         this.executor = connection.getExecutor();
         this.outgoingHandler = connection;
         this.connection.getIOState().addListener(this);
-        this.policy = policy;
+        this.policy = connection.getPolicy();
 
         addBean(this.connection);
     }

jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketSessionFactory.java

@@ -21,7 +21,6 @@ package org.eclipse.jetty.websocket.common;
 import java.net.URI;
 
 import org.eclipse.jetty.websocket.api.WebSocketConnectionListener;
-import org.eclipse.jetty.websocket.api.WebSocketPolicy;
 import org.eclipse.jetty.websocket.api.annotations.WebSocket;
 import org.eclipse.jetty.websocket.common.scopes.WebSocketContainerScope;
 
@@ -44,8 +43,8 @@ public class WebSocketSessionFactory implements SessionFactory
     }
 
     @Override
-    public WebSocketSession createSession(URI requestURI, Object websocket, WebSocketPolicy policy, LogicalConnection connection)
+    public WebSocketSession createSession(URI requestURI, Object websocket, LogicalConnection connection)
     {
-        return new WebSocketSession(containerScope, requestURI, websocket, policy, connection);
+        return new WebSocketSession(containerScope, requestURI, websocket, connection);
     }
 }

jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/io/AbstractWebSocketConnection.java

@@ -214,6 +214,7 @@ public abstract class AbstractWebSocketConnection extends AbstractConnection imp
     private final Scheduler scheduler;
     private final Generator generator;
     private final Parser parser;
+    private final WebSocketPolicy policy;
     private final WebSocketBehavior behavior;
     private final AtomicBoolean suspendToken;
     private final FrameFlusher flusher;
@@ -233,6 +234,7 @@ public abstract class AbstractWebSocketConnection extends AbstractConnection imp
                 endp.getLocalAddress().getPort(),
                 endp.getRemoteAddress().getAddress().getHostAddress(),
                 endp.getRemoteAddress().getPort());
+        this.policy = policy;
         this.behavior = policy.getBehavior();
         this.bufferPool = bufferPool;
         this.generator = new Generator(policy,bufferPool);
@@ -389,7 +391,12 @@ public abstract class AbstractWebSocketConnection extends AbstractConnection imp
     {
         return parser;
     }
-
+    
+    public WebSocketPolicy getPolicy()
+    {
+        return policy;
+    }
+    
     @Override
     public InetSocketAddress getRemoteAddress()
     {

jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/message/StringMessageSink.java

@@ -30,7 +30,7 @@ import org.eclipse.jetty.websocket.api.WebSocketPolicy;
 public class StringMessageSink implements MessageSink
 {
     private static final Logger LOG = Log.getLogger(StringMessageSink.class);
-    private final WebSocketPolicy policy;
+    private WebSocketPolicy policy;
     private final Function<String, Void> onMessageFunction;
     private Utf8StringBuilder utf;
     private int size = 0;

jetty-websocket/websocket-common/src/test/java/org/eclipse/jetty/websocket/common/io/LocalWebSocketConnection.java

@@ -36,6 +36,7 @@ import org.eclipse.jetty.websocket.common.CloseInfo;
 import org.eclipse.jetty.websocket.common.ConnectionState;
 import org.eclipse.jetty.websocket.common.LogicalConnection;
 import org.eclipse.jetty.websocket.common.io.IOState.ConnectionStateListener;
+import org.eclipse.jetty.websocket.common.scopes.WebSocketContainerScope;
 import org.junit.rules.TestName;
 
 public class LocalWebSocketConnection implements LogicalConnection, IncomingFrames, ConnectionStateListener
@@ -44,7 +45,7 @@ public class LocalWebSocketConnection implements LogicalConnection, IncomingFram
     private final String id;
     private final ByteBufferPool bufferPool;
     private final Executor executor;
-    private WebSocketPolicy policy = WebSocketPolicy.newServerPolicy();
+    private WebSocketPolicy policy;
     private IncomingFrames incoming;
     private IOState ioState = new IOState();
 
@@ -59,6 +60,7 @@ public class LocalWebSocketConnection implements LogicalConnection, IncomingFram
         this.bufferPool = bufferPool;
         this.executor = new ExecutorThreadPool();
         this.ioState.addListener(this);
+        this.policy = WebSocketPolicy.newServerPolicy();
     }
 
     public LocalWebSocketConnection(TestName testname, ByteBufferPool bufferPool)
@@ -66,6 +68,12 @@ public class LocalWebSocketConnection implements LogicalConnection, IncomingFram
         this(testname.getMethodName(),bufferPool);
     }
     
+    public LocalWebSocketConnection(TestName testname, WebSocketContainerScope containerScope)
+    {
+        this(testname.getMethodName(), containerScope.getBufferPool());
+        this.policy = containerScope.getPolicy();
+    }
+    
     @Override
     public Executor getExecutor()
     {
@@ -141,7 +149,13 @@ public class LocalWebSocketConnection implements LogicalConnection, IncomingFram
     {
         return 0;
     }
-
+    
+    @Override
+    public WebSocketPolicy getPolicy()
+    {
+        return policy;
+    }
+    
     @Override
     public InetSocketAddress getRemoteAddress()
     {

jetty-websocket/websocket-common/src/test/java/org/eclipse/jetty/websocket/common/io/LocalWebSocketSession.java

@@ -34,8 +34,7 @@ public class LocalWebSocketSession extends WebSocketSession
     public LocalWebSocketSession(WebSocketContainerScope containerScope, TestName testname, Object websocket)
     {
         super(containerScope,URI.create("ws://localhost/LocalWebSocketSesssion/" + testname.getMethodName()),websocket,
-                containerScope.getPolicy(),
+                new LocalWebSocketConnection(testname,containerScope));
-                new LocalWebSocketConnection(testname,containerScope.getBufferPool()));
         this.id = testname.getMethodName();
         outgoingCapture = new OutgoingFramesCapture();
         setOutgoingHandler(outgoingCapture);

jetty-websocket/websocket-common/src/test/java/org/eclipse/jetty/websocket/common/message/MessageWriterTest.java

@@ -82,7 +82,7 @@ public class MessageWriterTest
         remoteSocket = new TrackingSocket("remote");
         URI remoteURI = new URI("ws://localhost/remote");
         LocalWebSocketConnection remoteConnection = new LocalWebSocketConnection(bufferPool);
-        remoteSession = new WebSocketSession(containerScope,remoteURI,remoteSocket,policy,remoteConnection);
+        remoteSession = new WebSocketSession(containerScope,remoteURI,remoteSocket,remoteConnection);
         OutgoingFrames socketPipe = FramePipes.to(remoteSession);
         remoteSession.start();
         remoteSession.open();
@@ -91,7 +91,7 @@ public class MessageWriterTest
         TrackingSocket localSocket = new TrackingSocket("local");
         URI localURI = new URI("ws://localhost/local");
         LocalWebSocketConnection localConnection = new LocalWebSocketConnection(bufferPool);
-        session = new WebSocketSession(containerScope,localURI,localSocket,policy,localConnection);
+        session = new WebSocketSession(containerScope,localURI,localSocket,localConnection);
 
         // talk to our remote socket
         session.setOutgoingHandler(socketPipe);

jetty-websocket/websocket-common/src/test/java/org/eclipse/jetty/websocket/common/test/DummyConnection.java

@@ -108,7 +108,13 @@ public class DummyConnection implements LogicalConnection
     {
         return 0;
     }
-
+    
+    @Override
+    public WebSocketPolicy getPolicy()
+    {
+        return policy;
+    }
+    
     @Override
     public InetSocketAddress getRemoteAddress()
     {

jetty-websocket/websocket-server/src/main/java/org/eclipse/jetty/websocket/server/WebSocketServerFactory.java

@@ -265,7 +265,7 @@ public class WebSocketServerFactory extends ContainerLifeCycle implements WebSoc
             {
                 try
                 {
-                    return impl.createSession(requestURI, websocket, containerPolicy, connection);
+                    return impl.createSession(requestURI, websocket, connection);
                 }
                 catch (Throwable e)
                 {
@@ -562,7 +562,7 @@ public class WebSocketServerFactory extends ContainerLifeCycle implements WebSoc
         // Setup websocket connection
         AbstractWebSocketConnection wsConnection = new WebSocketServerConnection(endp, executor, scheduler, getPolicy().clonePolicy(), bufferPool);
 
-        extensionStack.setPolicy(containerPolicy);
+        extensionStack.setPolicy(wsConnection.getPolicy());
         extensionStack.configure(wsConnection.getParser());
         extensionStack.configure(wsConnection.getGenerator());