Bug 1105194 - Bail early if Websocket connection is being disconnected. r=smaug
authorAndrea Marchesini <amarchesini@mozilla.com>
Fri, 12 Dec 2014 14:37:07 -0500
changeset 219561 6122b7c0dd2eb1e8826240df7be2b8d8421b756a
parent 219560 9fa6f8d6c5f338f0b3c440008c43621acdce7910
child 219562 6f1696df360b086ad2f123317f5bf821eaa96891
push id27967
push userryanvm@gmail.com
push dateMon, 15 Dec 2014 18:52:54 +0000
treeherdermozilla-central@5d6e0d038f95 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewerssmaug
bugs1105194
milestone37.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1105194 - Bail early if Websocket connection is being disconnected. r=smaug
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;