Bug 1105194 - Bail early if Websocket connection is being disconnected. r=smaug, a=lsblakk
authorAndrea Marchesini <amarchesini@mozilla.com>
Fri, 12 Dec 2014 14:37:07 -0500
changeset 242512 7a83256f3e2f01706c9f373c1df534f1437fe46d
parent 242511 548a7408f76b5603d24762793d9c77cdf8ec9730
child 242513 5590fdb108c5143cb03edbf400a3a782be0e9c4c
push id4311
push userraliiev@mozilla.com
push dateMon, 12 Jan 2015 19:37:41 +0000
treeherdermozilla-beta@150c9fed433b [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewerssmaug, lsblakk
bugs1105194
milestone36.0a2
Bug 1105194 - Bail early if Websocket connection is being disconnected. r=smaug, a=lsblakk
dom/base/WebSocket.cpp
dom/base/WebSocket.h
--- a/dom/base/WebSocket.cpp
+++ b/dom/base/WebSocket.cpp
@@ -75,29 +75,32 @@ public:
   NS_DECL_NSIREQUEST
   NS_DECL_THREADSAFE_ISUPPORTS
   NS_DECL_NSIEVENTTARGET
 
   explicit WebSocketImpl(WebSocket* aWebSocket)
   : mWebSocket(aWebSocket)
   , mOnCloseScheduled(false)
   , mFailed(false)
-  , mDisconnected(false)
+  , mDisconnectingOrDisconnected(false)
   , mCloseEventWasClean(false)
   , mCloseEventCode(nsIWebSocketChannel::CLOSE_ABNORMAL)
   , mScriptLine(0)
   , mInnerWindowID(0)
   , mWorkerPrivate(nullptr)
 #ifdef DEBUG
   , mHasFeatureRegistered(false)
 #endif
+  , mIsMainThread(true)
+  , mWorkerShuttingDown(false)
   {
     if (!NS_IsMainThread()) {
       mWorkerPrivate = GetCurrentThreadWorkerPrivate();
       MOZ_ASSERT(mWorkerPrivate);
+      mIsMainThread = false;
     }
   }
 
   void AssertIsOnTargetThread() const
   {
     MOZ_ASSERT(IsTargetThread());
   }
 
@@ -147,31 +150,31 @@ public:
   // Dispatch a runnable to the right thread.
   nsresult DispatchRunnable(nsIRunnable* aRunnable);
 
   nsresult UpdateURI();
 
   void AddRefObject();
   void ReleaseObject();
 
-  void RegisterFeature();
+  bool RegisterFeature();
   void UnregisterFeature();
 
   nsresult CancelInternal();
 
   nsRefPtr<WebSocket> mWebSocket;
 
   nsCOMPtr<nsIWebSocketChannel> mChannel;
 
   bool mSecure; // if true it is using SSL and the wss scheme,
                 // otherwise it is using the ws scheme with no SSL
 
   bool mOnCloseScheduled;
   bool mFailed;
-  bool mDisconnected;
+  bool mDisconnectingOrDisconnected;
 
   // Set attributes of DOM 'onclose' message
   bool      mCloseEventWasClean;
   nsString  mCloseEventReason;
   uint16_t  mCloseEventCode;
 
   nsCString mAsciiHost;  // hostname
   uint32_t  mPort;
@@ -213,21 +216,24 @@ public:
     MOZ_ASSERT(mWebSocket);
     MutexAutoLock lock(mWebSocket->mMutex);
     mHasFeatureRegistered = aValue;
   }
 #endif
 
   nsWeakPtr mWeakLoadGroup;
 
+  bool mIsMainThread;
+  bool mWorkerShuttingDown;
+
 private:
   ~WebSocketImpl()
   {
     // If we threw during Init we never called disconnect
-    if (!mDisconnected) {
+    if (!mDisconnectingOrDisconnected) {
       Disconnect();
     }
   }
 };
 
 NS_IMPL_ISUPPORTS(WebSocketImpl,
                   nsIInterfaceRequestor,
                   nsIWebSocketListener,
@@ -308,16 +314,18 @@ nsresult
 WebSocketImpl::PrintErrorOnConsole(const char *aBundleURI,
                                    const char16_t *aError,
                                    const char16_t **aFormatStrings,
                                    uint32_t aFormatStringsLen)
 {
   // This method must run on the main thread.
 
   if (!NS_IsMainThread()) {
+    MOZ_ASSERT(mWorkerPrivate);
+
     nsRefPtr<PrintErrorOnConsoleRunnable> runnable =
       new PrintErrorOnConsoleRunnable(this, aBundleURI, aError, aFormatStrings,
                                       aFormatStringsLen);
     runnable->Dispatch(mWorkerPrivate->GetJSContext());
     return runnable->ErrorCode();
   }
 
   nsresult rv;
@@ -399,24 +407,52 @@ private:
   // A raw pointer because this runnable is sync.
   WebSocketImpl* mImpl;
 
   uint16_t mReasonCode;
   const nsACString& mReasonString;
   nsresult mRv;
 };
 
+class MOZ_STACK_CLASS MaybeDisconnect
+{
+public:
+  MaybeDisconnect(WebSocketImpl* aImpl)
+    : mImpl(aImpl)
+  {
+  }
+
+  ~MaybeDisconnect()
+  {
+    if (mImpl->mWorkerShuttingDown) {
+      mImpl->Disconnect();
+    }
+  }
+
+private:
+  WebSocketImpl* mImpl;
+};
+
 } // anonymous namespace
 
 nsresult
 WebSocketImpl::CloseConnection(uint16_t aReasonCode,
                                const nsACString& aReasonString)
 {
   AssertIsOnTargetThread();
 
+  if (mDisconnectingOrDisconnected) {
+    return NS_OK;
+  }
+
+  // If this method is called because the worker is going away, we will not
+  // receive the OnStop() method and we have to disconnect the WebSocket and
+  // release the WorkerFeature.
+  MaybeDisconnect md(this);
+
   uint16_t readyState = mWebSocket->ReadyState();
   if (readyState == WebSocket::CLOSING ||
       readyState == WebSocket::CLOSED) {
     return NS_OK;
   }
 
   // The common case...
   if (mChannel) {
@@ -479,16 +515,20 @@ WebSocketImpl::ConsoleError()
 }
 
 void
 WebSocketImpl::FailConnection(uint16_t aReasonCode,
                               const nsACString& aReasonString)
 {
   AssertIsOnTargetThread();
 
+  if (mDisconnectingOrDisconnected) {
+    return;
+  }
+
   ConsoleError();
   mFailed = true;
   CloseConnection(aReasonCode, aReasonString);
 }
 
 namespace {
 
 class DisconnectInternalRunnable MOZ_FINAL : public WorkerMainThreadRunnable
@@ -510,51 +550,56 @@ private:
   WebSocketImpl* mImpl;
 };
 
 } // anonymous namespace
 
 nsresult
 WebSocketImpl::Disconnect()
 {
-  if (mDisconnected) {
+  if (mDisconnectingOrDisconnected) {
     return NS_OK;
   }
 
   AssertIsOnTargetThread();
 
+  // Disconnect can be called from some control event (such as Notify() of
+  // WorkerFeature). This will be schedulated before any other sync/async
+  // runnable. In order to prevent some double Disconnect() calls, we use this
+  // boolean.
+  mDisconnectingOrDisconnected = true;
+
   // DisconnectInternal touches observers and nsILoadGroup and it must run on
   // the main thread.
 
   if (NS_IsMainThread()) {
     DisconnectInternal();
   } else {
     nsRefPtr<DisconnectInternalRunnable> runnable =
       new DisconnectInternalRunnable(this);
     runnable->Dispatch(mWorkerPrivate->GetJSContext());
   }
 
   // DontKeepAliveAnyMore() can release the object. So hold a reference to this
   // until the end of the method.
   nsRefPtr<WebSocketImpl> kungfuDeathGrip = this;
 
-  if (mWorkerPrivate && mWorkerFeature) {
-    UnregisterFeature();
-  }
-
   nsCOMPtr<nsIThread> mainThread;
   if (NS_FAILED(NS_GetMainThread(getter_AddRefs(mainThread))) ||
       NS_FAILED(NS_ProxyRelease(mainThread, mChannel))) {
     NS_WARNING("Failed to proxy release of channel, leaking instead!");
   }
 
-  mDisconnected = true;
   mWebSocket->DontKeepAliveAnyMore();
   mWebSocket->mImpl = nullptr;
 
+  if (mWorkerPrivate && mWorkerFeature) {
+    UnregisterFeature();
+  }
+
   // We want to release the WebSocket in the correct thread.
   mWebSocket = nullptr;
 
   return NS_OK;
 }
 
 void
 WebSocketImpl::DisconnectInternal()
@@ -580,16 +625,20 @@ WebSocketImpl::DisconnectInternal()
 // WebSocketImpl::nsIWebSocketListener methods:
 //-----------------------------------------------------------------------------
 
 nsresult
 WebSocketImpl::DoOnMessageAvailable(const nsACString& aMsg, bool isBinary)
 {
   AssertIsOnTargetThread();
 
+  if (mDisconnectingOrDisconnected) {
+    return NS_OK;
+  }
+
   int16_t readyState = mWebSocket->ReadyState();
   if (readyState == WebSocket::CLOSED) {
     NS_ERROR("Received message after CLOSED");
     return NS_ERROR_UNEXPECTED;
   }
 
   if (readyState == WebSocket::OPEN) {
     // Dispatch New Message
@@ -607,32 +656,46 @@ WebSocketImpl::DoOnMessageAvailable(cons
   return NS_OK;
 }
 
 NS_IMETHODIMP
 WebSocketImpl::OnMessageAvailable(nsISupports* aContext,
                                   const nsACString& aMsg)
 {
   AssertIsOnTargetThread();
+
+  if (mDisconnectingOrDisconnected) {
+    return NS_OK;
+  }
+
   return DoOnMessageAvailable(aMsg, false);
 }
 
 NS_IMETHODIMP
 WebSocketImpl::OnBinaryMessageAvailable(nsISupports* aContext,
                                         const nsACString& aMsg)
 {
   AssertIsOnTargetThread();
+
+  if (mDisconnectingOrDisconnected) {
+    return NS_OK;
+  }
+
   return DoOnMessageAvailable(aMsg, true);
 }
 
 NS_IMETHODIMP
 WebSocketImpl::OnStart(nsISupports* aContext)
 {
   AssertIsOnTargetThread();
 
+  if (mDisconnectingOrDisconnected) {
+    return NS_OK;
+  }
+
   int16_t readyState = mWebSocket->ReadyState();
 
   // This is the only function that sets OPEN, and should be called only once
   MOZ_ASSERT(readyState != WebSocket::OPEN,
              "readyState already OPEN! OnStart called twice?");
 
   // Nothing to do if we've already closed/closing
   if (readyState != WebSocket::CONNECTING) {
@@ -666,16 +729,20 @@ WebSocketImpl::OnStart(nsISupports* aCon
   return NS_OK;
 }
 
 NS_IMETHODIMP
 WebSocketImpl::OnStop(nsISupports* aContext, nsresult aStatusCode)
 {
   AssertIsOnTargetThread();
 
+  if (mDisconnectingOrDisconnected) {
+    return NS_OK;
+  }
+
   // We can be CONNECTING here if connection failed.
   // We can be OPEN if we have encountered a fatal protocol error
   // We can be CLOSING if close() was called and/or server initiated close.
   MOZ_ASSERT(mWebSocket->ReadyState() != WebSocket::CLOSED,
              "Shouldn't already be CLOSED when OnStop called");
 
   // called by network stack, not JS, so can dispatch JS events synchronously
   return ScheduleConnectionCloseEvents(aContext, aStatusCode, true);
@@ -714,30 +781,38 @@ WebSocketImpl::ScheduleConnectionCloseEv
   return NS_OK;
 }
 
 NS_IMETHODIMP
 WebSocketImpl::OnAcknowledge(nsISupports *aContext, uint32_t aSize)
 {
   AssertIsOnTargetThread();
 
+  if (mDisconnectingOrDisconnected) {
+    return NS_OK;
+  }
+
   if (aSize > mWebSocket->mOutgoingBufferedAmount) {
     return NS_ERROR_UNEXPECTED;
   }
 
   mWebSocket->mOutgoingBufferedAmount -= aSize;
   return NS_OK;
 }
 
 NS_IMETHODIMP
 WebSocketImpl::OnServerClose(nsISupports *aContext, uint16_t aCode,
                              const nsACString &aReason)
 {
   AssertIsOnTargetThread();
 
+  if (mDisconnectingOrDisconnected) {
+    return NS_OK;
+  }
+
   int16_t readyState = mWebSocket->ReadyState();
 
   MOZ_ASSERT(readyState != WebSocket::CONNECTING,
              "Received server close before connected?");
   MOZ_ASSERT(readyState != WebSocket::CLOSED,
              "Received server close after already closed!");
 
   // store code/string for onclose DOM event
@@ -797,26 +872,26 @@ WebSocketImpl::GetInterface(const nsIID&
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // WebSocket
 ////////////////////////////////////////////////////////////////////////////////
 
 WebSocket::WebSocket(nsPIDOMWindow* aOwnerWindow)
   : DOMEventTargetHelper(aOwnerWindow)
-  , mWorkerPrivate(nullptr)
+  , mIsMainThread(true)
   , mKeepingAlive(false)
   , mCheckMustKeepAlive(true)
   , mOutgoingBufferedAmount(0)
   , mBinaryType(dom::BinaryType::Blob)
   , mMutex("WebSocketImpl::mMutex")
   , mReadyState(CONNECTING)
 {
   mImpl = new WebSocketImpl(this);
-  mWorkerPrivate = mImpl->mWorkerPrivate;
+  mIsMainThread = mImpl->mIsMainThread;
 }
 
 WebSocket::~WebSocket()
 {
 }
 
 JSObject*
 WebSocket::WrapObject(JSContext* cx)
@@ -1053,16 +1128,23 @@ WebSocket::Constructor(const GlobalObjec
   nsRefPtr<WebSocketImpl> kungfuDeathGrip = webSocket->mImpl;
 
   bool connectionFailed = true;
 
   if (NS_IsMainThread()) {
     webSocket->mImpl->Init(aGlobal.Context(), principal, aUrl, protocolArray,
                            EmptyCString(), 0, aRv, &connectionFailed);
   } else {
+    // In workers we have to keep the worker alive using a feature in order to
+    // dispatch messages correctly.
+    if (!webSocket->mImpl->RegisterFeature()) {
+      aRv.Throw(NS_ERROR_FAILURE);
+      return nullptr;
+    }
+
     unsigned lineno;
     JS::AutoFilename file;
     if (!JS::DescribeScriptedCaller(aGlobal.Context(), &file, &lineno)) {
       NS_WARNING("Failed to get line number and filename in workers.");
     }
 
     nsRefPtr<InitRunnable> runnable =
       new InitRunnable(webSocket->mImpl, aUrl, protocolArray,
@@ -1070,34 +1152,35 @@ WebSocket::Constructor(const GlobalObjec
                        &connectionFailed);
     runnable->Dispatch(aGlobal.Context());
   }
 
   if (NS_WARN_IF(aRv.Failed())) {
     return nullptr;
   }
 
+  // It can be that we have been already disconnected because the WebSocket is
+  // gone away while we where initializing the webSocket.
+  if (!webSocket->mImpl) {
+    aRv.Throw(NS_ERROR_FAILURE);
+    return nullptr;
+  }
+
   // We don't return an error if the connection just failed. Instead we dispatch
   // an event.
   if (connectionFailed) {
     webSocket->mImpl->FailConnection(nsIWebSocketChannel::CLOSE_ABNORMAL);
   }
 
   // If we don't have a channel, the connection is failed and onerror() will be
   // called asynchrounsly.
   if (!webSocket->mImpl->mChannel) {
     return webSocket.forget();
   }
 
-  if (webSocket->mWorkerPrivate) {
-    // In workers we have to keep the worker alive using a feature in order to
-    // dispatch messages correctly.
-    webSocket->mImpl->RegisterFeature();
-  }
-
   class MOZ_STACK_CLASS ClearWebSocket
   {
   public:
     explicit ClearWebSocket(WebSocketImpl* aWebSocketImpl)
       : mWebSocketImpl(aWebSocketImpl)
       , mDone(false)
     {
     }
@@ -1135,16 +1218,23 @@ WebSocket::Constructor(const GlobalObjec
       new AsyncOpenRunnable(webSocket->mImpl, aRv);
     runnable->Dispatch(aGlobal.Context());
   }
 
   if (NS_WARN_IF(aRv.Failed())) {
     return nullptr;
   }
 
+  // It can be that we have been already disconnected because the WebSocket is
+  // gone away while we where initializing the webSocket.
+  if (!webSocket->mImpl) {
+    aRv.Throw(NS_ERROR_FAILURE);
+    return nullptr;
+  }
+
   cws.Done();
   return webSocket.forget();
 }
 
 NS_IMPL_CYCLE_COLLECTION_CLASS(WebSocket)
 
 NS_IMPL_CYCLE_COLLECTION_CAN_SKIP_BEGIN(WebSocket)
   bool isBlack = tmp->IsBlack();
@@ -1494,16 +1584,21 @@ WebSocketImpl::InitializeConnection()
 
   return NS_OK;
 }
 
 void
 WebSocketImpl::DispatchConnectionCloseEvents()
 {
   AssertIsOnTargetThread();
+
+  if (mDisconnectingOrDisconnected) {
+    return;
+  }
+
   mWebSocket->SetReadyState(WebSocket::CLOSED);
 
   // Call 'onerror' if needed
   if (mFailed) {
     nsresult rv =
       mWebSocket->CreateAndDispatchSimpleEvent(NS_LITERAL_STRING("error"));
     if (NS_FAILED(rv)) {
       NS_WARNING("Failed to dispatch the error event");
@@ -1554,18 +1649,19 @@ WebSocket::CreateAndDispatchMessageEvent
 
   AutoJSAPI jsapi;
 
   if (NS_IsMainThread()) {
     if (NS_WARN_IF(!jsapi.Init(GetOwner()))) {
       return NS_ERROR_FAILURE;
     }
   } else {
-    MOZ_ASSERT(mWorkerPrivate);
-    if (NS_WARN_IF(!jsapi.Init(mWorkerPrivate->GlobalScope()))) {
+    MOZ_ASSERT(!mIsMainThread);
+    MOZ_ASSERT(mImpl->mWorkerPrivate);
+    if (NS_WARN_IF(!jsapi.Init(mImpl->mWorkerPrivate->GlobalScope()))) {
       return NS_ERROR_FAILURE;
     }
   }
 
   return CreateAndDispatchMessageEvent(jsapi.cx(), aData, aIsBinary);
 }
 
 nsresult
@@ -1791,28 +1887,26 @@ WebSocketImpl::ParseURL(const nsAString&
 //      ("strong event listeners");
 //   2. there are outgoing not sent messages.
 //-----------------------------------------------------------------------------
 
 void
 WebSocket::UpdateMustKeepAlive()
 {
   // Here we could not have mImpl.
-  MOZ_ASSERT(NS_IsMainThread() == !mWorkerPrivate);
+  MOZ_ASSERT(NS_IsMainThread() == mIsMainThread);
 
   if (!mCheckMustKeepAlive || !mImpl) {
     return;
   }
 
   bool shouldKeepAlive = false;
   uint16_t readyState = ReadyState();
 
-  if (mWorkerPrivate && readyState != CLOSED) {
-    shouldKeepAlive = true;
-  } else if (mListenerManager) {
+  if (mListenerManager) {
     switch (readyState)
     {
       case CONNECTING:
       {
         if (mListenerManager->HasListenersFor(nsGkAtoms::onopen) ||
             mListenerManager->HasListenersFor(nsGkAtoms::onmessage) ||
             mListenerManager->HasListenersFor(nsGkAtoms::onerror) ||
             mListenerManager->HasListenersFor(nsGkAtoms::onclose)) {
@@ -1848,17 +1942,17 @@ WebSocket::UpdateMustKeepAlive()
     mImpl->AddRefObject();
   }
 }
 
 void
 WebSocket::DontKeepAliveAnyMore()
 {
   // Here we could not have mImpl.
-  MOZ_ASSERT(NS_IsMainThread() == !mWorkerPrivate);
+  MOZ_ASSERT(NS_IsMainThread() == mIsMainThread);
 
   if (mKeepingAlive) {
     MOZ_ASSERT(mImpl);
 
     mKeepingAlive = false;
     mImpl->ReleaseObject();
   }
 
@@ -1875,86 +1969,83 @@ public:
   {
   }
 
   bool Notify(JSContext* aCx, Status aStatus) MOZ_OVERRIDE
   {
     MOZ_ASSERT(aStatus > workers::Running);
 
     if (aStatus >= Canceling) {
+      mWebSocketImpl->mWorkerShuttingDown = true;
       mWebSocketImpl->CloseConnection(nsIWebSocketChannel::CLOSE_GOING_AWAY);
     }
 
     return true;
   }
 
   bool Suspend(JSContext* aCx)
   {
+    mWebSocketImpl->mWorkerShuttingDown = true;
     mWebSocketImpl->CloseConnection(nsIWebSocketChannel::CLOSE_GOING_AWAY);
     return true;
   }
 
 private:
   WebSocketImpl* mWebSocketImpl;
 };
 
 } // anonymous namespace
 
 void
 WebSocketImpl::AddRefObject()
 {
   AssertIsOnTargetThread();
   AddRef();
-
-  if (mWorkerPrivate && !mWorkerFeature) {
-    RegisterFeature();
-  }
 }
 
 void
 WebSocketImpl::ReleaseObject()
 {
   AssertIsOnTargetThread();
-
-  if (mWorkerPrivate && mWorkerFeature) {
-    UnregisterFeature();
-  }
-
   Release();
 }
 
-void
+bool
 WebSocketImpl::RegisterFeature()
 {
   mWorkerPrivate->AssertIsOnWorkerThread();
   MOZ_ASSERT(!mWorkerFeature);
   mWorkerFeature = new WebSocketWorkerFeature(this);
 
   JSContext* cx = GetCurrentThreadJSContext();
   if (!mWorkerPrivate->AddFeature(cx, mWorkerFeature)) {
     NS_WARNING("Failed to register a feature.");
     mWorkerFeature = nullptr;
-    return;
+    return false;
   }
 
 #ifdef DEBUG
   SetHasFeatureRegistered(true);
 #endif
+
+  return true;
 }
 
 void
 WebSocketImpl::UnregisterFeature()
 {
+  MOZ_ASSERT(mDisconnectingOrDisconnected);
   MOZ_ASSERT(mWorkerPrivate);
   mWorkerPrivate->AssertIsOnWorkerThread();
   MOZ_ASSERT(mWorkerFeature);
 
   JSContext* cx = GetCurrentThreadJSContext();
   mWorkerPrivate->RemoveFeature(cx, mWorkerFeature);
   mWorkerFeature = nullptr;
+  mWorkerPrivate = nullptr;
 
 #ifdef DEBUG
   SetHasFeatureRegistered(false);
 #endif
 }
 
 nsresult
 WebSocketImpl::UpdateURI()
@@ -2328,17 +2419,18 @@ private:
 } // anonymous namespace
 
 // Window closed, stop/reload button pressed, user navigated away from page, etc.
 NS_IMETHODIMP
 WebSocketImpl::Cancel(nsresult aStatus)
 {
   AssertIsOnMainThread();
 
-  if (mWorkerPrivate) {
+  if (!mIsMainThread) {
+    MOZ_ASSERT(mWorkerPrivate);
     nsRefPtr<CancelRunnable> runnable =
       new CancelRunnable(mWorkerPrivate, this);
     if (!runnable->Dispatch(nullptr)) {
       return NS_ERROR_FAILURE;
     }
 
     return NS_OK;
   }
@@ -2348,17 +2440,17 @@ WebSocketImpl::Cancel(nsresult aStatus)
 
 nsresult
 WebSocketImpl::CancelInternal()
 {
   AssertIsOnTargetThread();
 
    // If CancelInternal is called by a runnable, we may already be disconnected
    // by the time it runs.
-  if (mDisconnected) {
+  if (mDisconnectingOrDisconnected) {
     return NS_OK;
   }
 
   int64_t readyState = mWebSocket->ReadyState();
   if (readyState == WebSocket::CLOSING || readyState == WebSocket::CLOSED) {
     return NS_OK;
   }
 
@@ -2383,29 +2475,31 @@ WebSocketImpl::Resume()
 
 NS_IMETHODIMP
 WebSocketImpl::GetLoadGroup(nsILoadGroup** aLoadGroup)
 {
   AssertIsOnMainThread();
 
   *aLoadGroup = nullptr;
 
-  if (!mWorkerPrivate) {
+  if (mIsMainThread) {
     nsresult rv;
     nsIScriptContext* sc = mWebSocket->GetContextForEventHandlers(&rv);
     nsCOMPtr<nsIDocument> doc =
       nsContentUtils::GetDocumentFromScriptContext(sc);
 
     if (doc) {
       *aLoadGroup = doc->GetDocumentLoadGroup().take();
     }
 
     return NS_OK;
   }
 
+  MOZ_ASSERT(mWorkerPrivate);
+
   // Walk up to our containing page
   WorkerPrivate* wp = mWorkerPrivate;
   while (wp->GetParent()) {
     wp = wp->GetParent();
   }
 
   nsPIDOMWindow* window = wp->GetWindow();
   if (!window) {
@@ -2485,20 +2579,32 @@ private:
 };
 
 } // anonymous namespace
 
 NS_IMETHODIMP
 WebSocketImpl::Dispatch(nsIRunnable* aEvent, uint32_t aFlags)
 {
   // If the target is the main-thread we can just dispatch the runnable.
-  if (!mWorkerPrivate) {
+  if (mIsMainThread) {
     return NS_DispatchToMainThread(aEvent);
   }
 
+  // No messages when disconnected.
+  if (mDisconnectingOrDisconnected) {
+    NS_WARNING("Dispatching a WebSocket event after the disconnection!");
+    return NS_OK;
+  }
+
+  if (mWorkerShuttingDown) {
+    return NS_OK;
+  }
+
+  MOZ_ASSERT(mWorkerPrivate);
+
 #ifdef DEBUG
   MOZ_ASSERT(HasFeatureRegistered());
 #endif
 
   // If the target is a worker, we have to use a custom WorkerRunnableDispatcher
   // runnable.
   nsRefPtr<WorkerRunnableDispatcher> event =
     new WorkerRunnableDispatcher(mWorkerPrivate, aEvent);
@@ -2514,19 +2620,19 @@ WebSocketImpl::IsOnCurrentThread(bool* a
 {
   *aResult = IsTargetThread();
   return NS_OK;
 }
 
 bool
 WebSocketImpl::IsTargetThread() const
 {
-  return NS_IsMainThread() == !mWorkerPrivate;
+  return NS_IsMainThread() == mIsMainThread;
 }
 
 void
 WebSocket::AssertIsOnTargetThread() const
 {
-  MOZ_ASSERT(NS_IsMainThread() == !mWorkerPrivate);
+  MOZ_ASSERT(NS_IsMainThread() == mIsMainThread);
 }
 
 } // dom namespace
 } // mozilla namespace
--- a/dom/base/WebSocket.h
+++ b/dom/base/WebSocket.h
@@ -168,17 +168,17 @@ private:
             ErrorResult& aRv);
 
   void AssertIsOnTargetThread() const;
 
   // Raw pointer because this WebSocketImpl is created, managed and destroyed by
   // WebSocket.
   WebSocketImpl* mImpl;
 
-  workers::WorkerPrivate* mWorkerPrivate;
+  bool mIsMainThread;
 
   bool mKeepingAlive;
   bool mCheckMustKeepAlive;
 
   uint32_t mOutgoingBufferedAmount;
 
   // related to the WebSocket constructor steps
   nsString mOriginalURL;