From b95e454f7d9208639b91bf1b134046476d7d1803 Mon Sep 17 00:00:00 2001
From: Joakim Erdfelt <joakim.erdfelt@gmail.com>
Date: Tue, 11 Mar 2014 12:12:44 -0700
Subject: [PATCH] 417022 - Access current HttpConnection from Request not
 ThreadLocal

+ Adjusting WebSocketServerFactory.upgrade() to receive HttpConnection
  from Request object instead of ThreadLocal.
---
 .../websocket/jsr356/server/MemoryUsageTest.java  |  4 +++-
 .../websocket/server/WebSocketServerFactory.java  | 15 ++++++++++-----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/MemoryUsageTest.java b/jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/MemoryUsageTest.java
index 29c5fee44c9..d5d4ba9aaf2 100644
--- a/jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/MemoryUsageTest.java
+++ b/jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/MemoryUsageTest.java
@@ -18,6 +18,8 @@
 
 package org.eclipse.jetty.websocket.jsr356.server;
 
+import static org.hamcrest.Matchers.*;
+
 import java.lang.management.ManagementFactory;
 import java.lang.management.MemoryMXBean;
 import java.lang.management.MemoryUsage;
@@ -116,7 +118,7 @@ public class MemoryUsageTest
 
         // Assume no more than 25 KiB per session pair (client and server).
         long expected = 25 * 1024 * sessionCount;
-        Assert.assertTrue(heapUsed < expected);
+        Assert.assertThat("heap used", heapUsed,lessThan(expected));
     }
 
     private static abstract class EndpointAdapter extends Endpoint implements MessageHandler.Whole<String>
diff --git a/jetty-websocket/websocket-server/src/main/java/org/eclipse/jetty/websocket/server/WebSocketServerFactory.java b/jetty-websocket/websocket-server/src/main/java/org/eclipse/jetty/websocket/server/WebSocketServerFactory.java
index 9609ce2506e..bebe7c0c379 100644
--- a/jetty-websocket/websocket-server/src/main/java/org/eclipse/jetty/websocket/server/WebSocketServerFactory.java
+++ b/jetty-websocket/websocket-server/src/main/java/org/eclipse/jetty/websocket/server/WebSocketServerFactory.java
@@ -164,9 +164,12 @@ public class WebSocketServerFactory extends ContainerLifeCycle implements WebSoc
                 return false;
             }
 
+            // Get the original HTTPConnection
+            HttpConnection connection = (HttpConnection)request.getAttribute("org.eclipse.jetty.server.HttpConnection");
+            
             // Send the upgrade
             EventDriver driver = eventDriverFactory.wrap(websocketPojo);
-            return upgrade(sockreq, sockresp, driver);
+            return upgrade(connection, sockreq, sockresp, driver);
         }
         catch (URISyntaxException e)
         {
@@ -407,12 +410,13 @@ public class WebSocketServerFactory extends ContainerLifeCycle implements WebSoc
      * This method will not normally return, but will instead throw a UpgradeConnectionException, to exit HTTP handling and initiate WebSocket handling of the
      * connection.
      *
+     * @param http     the raw http connection
      * @param request  The request to upgrade
      * @param response The response to upgrade
      * @param driver   The websocket handler implementation to use
      * @throws IOException
      */
-    public boolean upgrade(ServletUpgradeRequest request, ServletUpgradeResponse response, EventDriver driver) throws IOException
+    private boolean upgrade(HttpConnection http, ServletUpgradeRequest request, ServletUpgradeResponse response, EventDriver driver) throws IOException
     {
         if (!"websocket".equalsIgnoreCase(request.getHeader("Upgrade")))
         {
@@ -477,11 +481,12 @@ public class WebSocketServerFactory extends ContainerLifeCycle implements WebSoc
             extensionStack.negotiate(request.getExtensions());
         }
 
-        // Create connection
-        HttpConnection http = HttpConnection.getCurrentConnection();
+        // Get original HTTP connection
         EndPoint endp = http.getEndPoint();
         Executor executor = http.getConnector().getExecutor();
         ByteBufferPool bufferPool = http.getConnector().getByteBufferPool();
+        
+        // Setup websocket connection
         WebSocketServerConnection wsConnection = new WebSocketServerConnection(endp, executor, scheduler, driver.getPolicy(), bufferPool);
 
         extensionStack.setPolicy(driver.getPolicy());
@@ -526,7 +531,7 @@ public class WebSocketServerFactory extends ContainerLifeCycle implements WebSoc
             throw new IOException("Unable to start Extension Stack", e);
         }
 
-        // Tell jetty about the new connection
+        // Tell jetty about the new upgraded connection
         request.setServletAttribute(HttpConnection.UPGRADE_CONNECTION_ATTRIBUTE, wsConnection);
 
         // Process (version specific) handshake response