bug 712572 part 2 r=jduell

netwerk/protocol/websocket/WebSocketChannel.cpp

@@ -382,7 +382,7 @@ NS_IMPL_THREADSAFE_ISUPPORTS1(OutboundEnqueuer, nsIRunnable)
 class nsWSAdmissionManager
 {
 public:
-  nsWSAdmissionManager() : mConnectedCount(0)
+  nsWSAdmissionManager() : mSessionCount(0)
   {
     MOZ_COUNT_CTOR(nsWSAdmissionManager);
   }
@@ -500,19 +500,19 @@ public:
     return false;
   }
 
-  void IncrementConnectedCount()
+  void IncrementSessionCount()
   {
-    PR_ATOMIC_INCREMENT(&mConnectedCount);
+    PR_ATOMIC_INCREMENT(&mSessionCount);
   }
 
-  void DecrementConnectedCount()
+  void DecrementSessionCount()
   {
-    PR_ATOMIC_DECREMENT(&mConnectedCount);
+    PR_ATOMIC_DECREMENT(&mSessionCount);
   }
 
-  PRInt32 ConnectedCount()
+  PRInt32 SessionCount()
   {
-    return mConnectedCount;
+    return mSessionCount;
   }
 
 private:
@@ -534,9 +534,9 @@ private:
     return -1;
   }
 
-  // ConnectedCount might be decremented from the main or the socket
+  // SessionCount might be decremented from the main or the socket
   // thread, so manage it with atomic counters
-  PRInt32 mConnectedCount;
+  PRInt32 mSessionCount;
 };
 
 //-----------------------------------------------------------------------------
@@ -714,6 +714,9 @@ WebSocketChannel::WebSocketChannel() :
   if (!sWebSocketAdmissions)
     sWebSocketAdmissions = new nsWSAdmissionManager();
 
+  // The active session limit is enforced in AsyncOpen()
+  sWebSocketAdmissions->IncrementSessionCount();
+
   mFramePtr = mBuffer = static_cast<PRUint8 *>(moz_xmalloc(mBufferSize));
 }
 
@@ -721,6 +724,9 @@ WebSocketChannel::~WebSocketChannel()
 {
   LOG(("WebSocketChannel::~WebSocketChannel() %p\n", this));
 
+  if (sWebSocketAdmissions)
+    sWebSocketAdmissions->DecrementSessionCount();
+
   // this stop is a nop if the normal connect/close is followed
   mStopped = 1;
   StopSession(NS_ERROR_UNEXPECTED);
@@ -1534,8 +1540,6 @@ WebSocketChannel::CleanupConnection()
   }
 
   if (mSocketIn) {
-    if (sWebSocketAdmissions)
-      sWebSocketAdmissions->DecrementConnectedCount();
     mSocketIn->AsyncWait(nsnull, 0, 0, nsnull);
     mSocketIn = nsnull;
   }
@@ -1612,7 +1616,7 @@ WebSocketChannel::StopSession(nsresult reason)
   }
 
   if (!mTCPClosed && mTransport && sWebSocketAdmissions &&
-    sWebSocketAdmissions->ConnectedCount() < kLingeringCloseThreshold) {
+      sWebSocketAdmissions->SessionCount() < kLingeringCloseThreshold) {
 
     // 7.1.1 says that the client SHOULD wait for the server to close the TCP
     // connection. This is so we can reuse port numbers before 2 MSL expires,
@@ -1885,14 +1889,6 @@ WebSocketChannel::StartWebsocketData()
 {
   LOG(("WebSocketChannel::StartWebsocketData() %p", this));
 
-  if (sWebSocketAdmissions &&
-    sWebSocketAdmissions->ConnectedCount() > mMaxConcurrentConnections) {
-    LOG(("WebSocketChannel max concurrency %d exceeded "
-         "in OnTransportAvailable()", mMaxConcurrentConnections));
-    AbortSession(NS_ERROR_SOCKET_CREATE_FAILED);
-    return NS_OK;
-  }
-
   return mSocketIn->AsyncWait(this, 0, 0, mSocketThread);
 }
 
@@ -1910,8 +1906,10 @@ WebSocketChannel::OnLookupComplete(nsICancelable *aRequest,
   NS_ABORT_IF_FALSE(aRequest == mDNSRequest || mStopped,
                     "wrong dns request");
 
-  if (mStopped)
+  if (mStopped) {
+    LOG(("WebSocketChannel::OnLookupComplete: Request Already Stopped\n"));
     return NS_OK;
+  }
 
   mDNSRequest = nsnull;
 
@@ -2223,15 +2221,16 @@ WebSocketChannel::AsyncOpen(nsIURI *aURI,
     }
   }
 
+  if (sWebSocketAdmissions)
+    LOG(("WebSocketChannel::AsyncOpen %p sessionCount=%d max=%d\n", this,
+         sWebSocketAdmissions->SessionCount(), mMaxConcurrentConnections));
+
   if (sWebSocketAdmissions &&
-      sWebSocketAdmissions->ConnectedCount() >= mMaxConcurrentConnections)
+      sWebSocketAdmissions->SessionCount() >= mMaxConcurrentConnections)
   {
-    // Checking this early creates an optimal fast-fail, but it is
-    // also a time-of-check-time-of-use problem. So we will check again
-    // after the handshake is complete to catch anything that sneaks
-    // through the race condition.
-    LOG(("WebSocketChannel: max concurrency %d exceeded",
-         mMaxConcurrentConnections));
+    LOG(("WebSocketChannel: max concurrency %d exceeded (%d)",
+         mMaxConcurrentConnections,
+         sWebSocketAdmissions->SessionCount()));
 
     // WebSocket connections are expected to be long lived, so return
     // an error here instead of queueing
@@ -2404,8 +2403,6 @@ WebSocketChannel::OnTransportAvailable(nsISocketTransport *aTransport,
   mTransport = aTransport;
   mSocketIn = aSocketIn;
   mSocketOut = aSocketOut;
-  if (sWebSocketAdmissions)
-    sWebSocketAdmissions->IncrementConnectedCount();
 
   nsresult rv;
   rv = mTransport->SetEventSink(nsnull, nsnull);