Issue #1442 - Part 9: Stream shutdown handling.

https://bugzilla.mozilla.org/show_bug.cgi?id=1128959

2 files changed, 182 insertions, 4 deletions

diff --git a/dom/fetch/FetchStream.cpp b/dom/fetch/FetchStream.cpp
index 7b165b2a9d..e57991fcc2 100644
--- a/dom/fetch/FetchStream.cpp
+++ b/dom/fetch/FetchStream.cpp
@@ -8,6 +8,9 @@
 #include "nsITransport.h"
 #include "nsIStreamTransportService.h"
 #include "nsProxyRelease.h"
+#include "WorkerPrivate.h"
+#include "WorkerRunnable.h"
+#include "Workers.h"
 
 #include "mozilla/dom/DOMError.h"
 
@@ -19,7 +22,67 @@ static NS_DEFINE_CID(kStreamTransportServiceCID,
 namespace mozilla {
 namespace dom {
 
-NS_IMPL_ISUPPORTS(FetchStream, nsIInputStreamCallback)
+using namespace workers;
+
+namespace {
+
+class FetchStreamWorkerHolder final : public WorkerHolder
+{
+public:
+  explicit FetchStreamWorkerHolder(FetchStream* aStream)
+    : WorkerHolder()
+    , mStream(aStream)
+    , mWasNotified(false)
+  {}
+
+  bool Notify(Status aStatus) override
+  {
+    if (!mWasNotified) {
+      mWasNotified = true;
+      mStream->Close();
+    }
+
+    return true;
+  }
+
+  WorkerPrivate* GetWorkerPrivate() const
+  {
+    return mWorkerPrivate;
+  }
+
+private:
+  RefPtr<FetchStream> mStream;
+  bool mWasNotified;
+};
+
+class FetchStreamWorkerHolderShutdown final : public WorkerControlRunnable
+{
+public:
+  FetchStreamWorkerHolderShutdown(WorkerPrivate* aWorkerPrivate,
+                                  UniquePtr<WorkerHolder>&& aHolder,
+                                  nsCOMPtr<nsIGlobalObject>&& aGlobal)
+    : WorkerControlRunnable(aWorkerPrivate)
+    , mHolder(Move(aHolder))
+    , mGlobal(Move(aGlobal))
+  {}
+
+  bool
+  WorkerRun(JSContext* aCx, WorkerPrivate* aWorkerPrivate) override
+  {
+    mHolder = nullptr;
+    mGlobal = nullptr;
+    return true;
+  }
+
+private:
+  UniquePtr<WorkerHolder> mHolder;
+  nsCOMPtr<nsIGlobalObject> mGlobal;
+};
+
+} // anonymous
+
+NS_IMPL_ISUPPORTS(FetchStream, nsIInputStreamCallback, nsIObserver,
+                  nsISupportsWeakReference)
 
 /* static */ JSObject*
 FetchStream::Create(JSContext* aCx, nsIGlobalObject* aGlobal,
@@ -30,6 +93,35 @@ FetchStream::Create(JSContext* aCx, nsIGlobalObject* aGlobal,
 
   RefPtr<FetchStream> stream = new FetchStream(aGlobal, aInputStream);
 
+  if (NS_IsMainThread()) {
+    nsCOMPtr<nsIObserverService> os = mozilla::services::GetObserverService();
+    if (NS_WARN_IF(!os)) {
+      aRv.Throw(NS_ERROR_FAILURE);
+      return nullptr;
+    }
+
+    aRv = os->AddObserver(stream, DOM_WINDOW_DESTROYED_TOPIC, true);
+    if (NS_WARN_IF(aRv.Failed())) {
+      return nullptr;
+    }
+
+  } else {
+    WorkerPrivate* workerPrivate = GetWorkerPrivateFromContext(aCx);
+    MOZ_ASSERT(workerPrivate);
+
+    UniquePtr<FetchStreamWorkerHolder> holder(
+      new FetchStreamWorkerHolder(stream));
+    if (NS_WARN_IF(!holder->HoldWorker(workerPrivate, Closing))) {
+      aRv.Throw(NS_ERROR_DOM_INVALID_STATE_ERR);
+      return nullptr;
+    }
+
+    // Note, this will create a ref-cycle between the holder and the stream.
+    // The cycle is broken when the stream is closed or the worker begins
+    // shutting down.
+    stream->mWorkerHolder = Move(holder);
+  }
+
   if (!JS::HasReadableStreamCallbacks(aCx)) {
     JS::SetReadableStreamCallbacks(aCx,
                                    &FetchStream::RequestDataCallback,
@@ -228,11 +320,16 @@ FetchStream::FinalizeCallback(void* aUnderlyingSource, uint8_t aFlags)
   MOZ_DIAGNOSTIC_ASSERT(aUnderlyingSource);
   MOZ_DIAGNOSTIC_ASSERT(aFlags == FETCH_STREAM_FLAG);
 
+  // This can be called in any thread.
+
   RefPtr<FetchStream> stream =
     dont_AddRef(static_cast<FetchStream*>(aUnderlyingSource));
 
-  stream->mState = eClosed;
-  stream->mReadableStream = nullptr;
+  if (stream->mState == eClosed) {
+    return;
+  }
+
+  stream->CloseAndReleaseObjects();
 }
 
 FetchStream::FetchStream(nsIGlobalObject* aGlobal,
@@ -248,7 +345,6 @@ FetchStream::FetchStream(nsIGlobalObject* aGlobal,
 
 FetchStream::~FetchStream()
 {
-  NS_ProxyRelease(mOwningEventTarget, mGlobal.forget());
 }
 
 void
@@ -332,5 +428,70 @@ FetchStream::OnInputStreamReady(nsIAsyncInputStream* aStream)
   return NS_OK;
 }
 
+void
+FetchStream::Close()
+{
+  if (mState == eClosed) {
+    return;
+  }
+
+  AutoJSAPI jsapi;
+  if (NS_WARN_IF(!jsapi.Init(mGlobal))) {
+    return;
+  }
+
+  JSContext* cx = jsapi.cx();
+  JS::Rooted<JSObject*> stream(cx, mReadableStream);
+  JS::ReadableStreamClose(cx, stream);
+
+  CloseAndReleaseObjects();
+}
+
+void
+FetchStream::CloseAndReleaseObjects()
+{
+  MOZ_DIAGNOSTIC_ASSERT(mState != eClosed);
+  mState = eClosed;
+
+  if (mWorkerHolder) {
+    RefPtr<FetchStreamWorkerHolderShutdown> r =
+      new FetchStreamWorkerHolderShutdown(
+        static_cast<FetchStreamWorkerHolder*>(mWorkerHolder.get())->GetWorkerPrivate(),
+        Move(mWorkerHolder), Move(mGlobal));
+    r->Dispatch();
+  } else {
+    RefPtr<FetchStream> self = this;
+    RefPtr<Runnable> r = NS_NewRunnableFunction(
+      [self] () {
+        nsCOMPtr<nsIObserverService> os = mozilla::services::GetObserverService();
+        if (os) {
+          os->RemoveObserver(self, DOM_WINDOW_DESTROYED_TOPIC);
+        }
+        self->mGlobal = nullptr;
+      });
+
+    NS_DispatchToMainThread(r);
+  }
+}
+
+// nsIObserver
+// -----------
+
+NS_IMETHODIMP
+FetchStream::Observe(nsISupports* aSubject, const char* aTopic,
+                     const char16_t* aData)
+{
+  AssertIsOnMainThread();
+
+  MOZ_ASSERT(strcmp(aTopic, DOM_WINDOW_DESTROYED_TOPIC) == 0);
+
+  nsCOMPtr<nsPIDOMWindowInner> window = do_QueryInterface(mGlobal);
+  if (SameCOMIdentity(aSubject, window)) {
+    Close();
+  }
+
+  return NS_OK;
+}
+
 } // dom namespace
 } // mozilla namespace
diff --git a/dom/fetch/FetchStream.h b/dom/fetch/FetchStream.h
index ebe94e205b..b18b7356ff 100644
--- a/dom/fetch/FetchStream.h
+++ b/dom/fetch/FetchStream.h
@@ -9,7 +9,9 @@
 #include "Fetch.h"
 #include "jsapi.h"
 #include "nsIAsyncInputStream.h"
+#include "nsIObserver.h"
 #include "nsISupportsImpl.h"
+#include "nsWeakReference.h"
 
 class nsIGlobalObject;
 
@@ -18,16 +20,26 @@ class nsIInputStream;
 namespace mozilla {
 namespace dom {
 
+namespace workers {
+class WorkerHolder;
+}
+
 class FetchStream final : public nsIInputStreamCallback
+                        , public nsIObserver
+                        , public nsSupportsWeakReference
 {
 public:
   NS_DECL_THREADSAFE_ISUPPORTS
   NS_DECL_NSIINPUTSTREAMCALLBACK
+  NS_DECL_NSIOBSERVER
 
   static JSObject*
   Create(JSContext* aCx, nsIGlobalObject* aGlobal,
          nsIInputStream* aInputStream, ErrorResult& aRv);
 
+  void
+  Close();
+
 private:
   FetchStream(nsIGlobalObject* aGlobal, nsIInputStream* aInputStream);
   ~FetchStream();
@@ -63,6 +75,9 @@ private:
   void
   ErrorPropagation(JSContext* aCx, JS::HandleObject aStream, nsresult aRv);
 
+  void
+  CloseAndReleaseObjects();
+
   // Common methods
 
   enum State {
@@ -98,6 +113,8 @@ private:
 
   nsCOMPtr<nsIEventTarget> mOwningEventTarget;
 
+  UniquePtr<workers::WorkerHolder> mWorkerHolder;
+
   JS::Heap<JSObject*> mReadableStream;
 };