From 58651be5dd900ce18a67d343eca4abc61955d653 Mon Sep 17 00:00:00 2001
From: Gijs Kruitbosch <gijskruitbosch@gmail.com>
Date: Wed, 2 Sep 2020 23:15:21 +0000
Subject: [PATCH] Bug 1658202 - move flushes caused by closing the stream away
 from the main thread, r=valentin

Depends on D88730

Differential Revision: https://phabricator.services.mozilla.com/D88731
---
 dom/webbrowserpersist/nsWebBrowserPersist.cpp | 35 ++++++++++++++++++-
 dom/webbrowserpersist/nsWebBrowserPersist.h   |  6 +++-
 .../components/downloads/DownloadLegacy.jsm   |  5 ---
 3 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/dom/webbrowserpersist/nsWebBrowserPersist.cpp b/dom/webbrowserpersist/nsWebBrowserPersist.cpp
index 2a63f0ddbe3c..1328c08a5578 100644
--- a/dom/webbrowserpersist/nsWebBrowserPersist.cpp
+++ b/dom/webbrowserpersist/nsWebBrowserPersist.cpp
@@ -889,8 +889,34 @@ NS_IMETHODIMP nsWebBrowserPersist::OnStopRequest(nsIRequest* request,
       SendErrorStatusChange(true, status, request, data->mFile);
     }
 
+    // If there is a stream ref and we weren't canceled,
+    // close it away from the main thread.
+    // We don't do this when there's an error/cancelation,
+    // because our consumer may try to delete the file, which will error
+    // if we're still holding on to it, so we have to close it pronto.
+    {
+      MutexAutoLock lock(data->mStreamMutex);
+      if (data->mStream && NS_SUCCEEDED(status) && !mCancel) {
+        if (!mBackgroundQueue) {
+          nsresult rv = NS_CreateBackgroundTaskQueue(
+              "WebBrowserPersist", getter_AddRefs(mBackgroundQueue));
+          if (NS_FAILED(rv)) {
+            return rv;
+          }
+        }
+        // Now steal the stream ref and close it away from the main thread,
+        // keeping the promise around so we don't finish before all files
+        // are flushed and closed.
+        mFileClosePromises.AppendElement(InvokeAsync(
+            mBackgroundQueue, __func__, [stream = std::move(data->mStream)]() {
+              nsresult rv = stream->Close();
+              // We don't care if closing failed; we don't care in the
+              // destructor either...
+              return ClosePromise::CreateAndResolve(rv, __func__);
+            }));
+      }
+    }
     MutexAutoLock lock(mOutputMapMutex);
-    // This will automatically close the output stream
     mOutputMap.Remove(keyPtr);
   } else {
     // if we didn't find the data in mOutputMap, try mUploadList
@@ -2275,7 +2301,14 @@ void nsWebBrowserPersist::EndDownload(nsresult aResult) {
   if (NS_SUCCEEDED(mPersistResult) && NS_FAILED(aResult)) {
     mPersistResult = aResult;
   }
+  ClosePromise::All(GetCurrentSerialEventTarget(), mFileClosePromises)
+      ->Then(GetCurrentSerialEventTarget(), __func__,
+             [self = RefPtr{this}, aResult]() {
+               self->EndDownloadInternal(aResult);
+             });
+}
 
+void nsWebBrowserPersist::EndDownloadInternal(nsresult aResult) {
   // mCompleted needs to be set before issuing the stop notification.
   // (Bug 1224437)
   mCompleted = true;
diff --git a/dom/webbrowserpersist/nsWebBrowserPersist.h b/dom/webbrowserpersist/nsWebBrowserPersist.h
index 034fe7153b4b..5a113f094cd7 100644
--- a/dom/webbrowserpersist/nsWebBrowserPersist.h
+++ b/dom/webbrowserpersist/nsWebBrowserPersist.h
@@ -31,6 +31,8 @@
 class nsIStorageStream;
 class nsIWebBrowserPersistDocument;
 
+using ClosePromise = mozilla::MozPromise<nsresult, nsresult, true>;
+
 class nsWebBrowserPersist final : public nsIInterfaceRequestor,
                                   public nsIWebBrowserPersist,
                                   public nsIStreamListener,
@@ -124,8 +126,9 @@ class nsWebBrowserPersist final : public nsIInterfaceRequestor,
 
   nsresult FixRedirectedChannelEntry(nsIChannel* aNewChannel);
 
-  void EndDownload(nsresult aResult);
   void FinishDownload();
+  void EndDownload(nsresult aResult);
+  void EndDownloadInternal(nsresult aResult);
   void SerializeNextFile();
   void CalcTotalProgress();
 
@@ -153,6 +156,7 @@ class nsWebBrowserPersist final : public nsIInterfaceRequestor,
   nsClassHashtable<nsISupportsHashKey, OutputData> mOutputMap;
   nsClassHashtable<nsISupportsHashKey, UploadData> mUploadList;
   nsCOMPtr<nsISerialEventTarget> mBackgroundQueue;
+  nsTArray<RefPtr<ClosePromise>> mFileClosePromises;
   nsClassHashtable<nsCStringHashKey, URIData> mURIMap;
   nsCOMPtr<nsIWebBrowserPersistURIMap> mFlatURIMap;
   nsTArray<mozilla::UniquePtr<WalkData>> mWalkStack;
diff --git a/toolkit/components/downloads/DownloadLegacy.jsm b/toolkit/components/downloads/DownloadLegacy.jsm
index ee6db66890d4..4cd1a627c49c 100644
--- a/toolkit/components/downloads/DownloadLegacy.jsm
+++ b/toolkit/components/downloads/DownloadLegacy.jsm
@@ -122,11 +122,6 @@ DownloadLegacyTransfer.prototype = {
             // Only cancel if the object executing the download is still running.
             if (this._cancelable && !this._componentFailed) {
               this._cancelable.cancel(Cr.NS_ERROR_ABORT);
-              if (this._cancelable instanceof Ci.nsIWebBrowserPersist) {
-                // This component will not send the STATE_STOP notification.
-                download.saver.onTransferFinished(Cr.NS_ERROR_ABORT);
-                this._cancelable = null;
-              }
             }
           });
         })