Bug 1896232 - Wait for screenshot to download before closing overlay. r=sfoster

Differential Revision: https://phabricator.services.mozilla.com/D216083

browser/actors/ScreenshotsComponentChild.sys.mjs

@@ -18,6 +18,7 @@ export class ScreenshotsComponentChild extends JSWindowActorChild {
   #scrollTask;
   #overlay;
   #preventableEventsAdded = false;
+  #listening = false;
 
   static OVERLAY_EVENTS = [
     "click",
@@ -172,7 +173,6 @@ export class ScreenshotsComponentChild extends JSWindowActorChild {
       closeOverlay: false,
       reason,
     });
-    this.endScreenshotsOverlay();
   }
 
   /**
@@ -182,7 +182,6 @@ export class ScreenshotsComponentChild extends JSWindowActorChild {
   requestCopyScreenshot(region) {
     region.devicePixelRatio = this.contentWindow.devicePixelRatio;
     this.sendAsyncMessage("Screenshots:CopyScreenshot", { region });
-    this.endScreenshotsOverlay({ doNotResetMethods: true });
   }
 
   /**
@@ -195,7 +194,6 @@ export class ScreenshotsComponentChild extends JSWindowActorChild {
       title: this.getDocumentTitle(),
       region,
     });
-    this.endScreenshotsOverlay({ doNotResetMethods: true });
   }
 
   getDocumentTitle() {
@@ -258,6 +256,7 @@ export class ScreenshotsComponentChild extends JSWindowActorChild {
     this.contentWindow.addEventListener("resize", this);
     this.contentWindow.addEventListener("scroll", this);
     this.addOverlayEventListeners();
+    this.#listening = true;
   }
 
   addOverlayEventListeners() {
@@ -291,20 +290,26 @@ export class ScreenshotsComponentChild extends JSWindowActorChild {
       return false;
     }
     await this.documentIsReady();
-    let overlay =
-      this.overlay ||
-      (this.#overlay = new lazy.ScreenshotsOverlay(this.document));
+    if (!this.overlay) {
+      this.#overlay = new lazy.ScreenshotsOverlay(this.document);
+    }
     this.addEventListeners();
 
-    overlay.initialize();
+    this.overlay.initialize();
     return true;
   }
 
   removeEventListeners() {
+    if (!this.#listening) {
+      return;
+    }
+
     this.contentWindow.removeEventListener("beforeunload", this);
     this.contentWindow.removeEventListener("resize", this);
     this.contentWindow.removeEventListener("scroll", this);
     this.removeOverlayEventListeners();
+
+    this.#listening = false;
   }
 
   removeOverlayEventListeners() {
@@ -336,8 +341,7 @@ export class ScreenshotsComponentChild extends JSWindowActorChild {
   }
 
   didDestroy() {
-    this.#resizeTask?.disarm();
-    this.#scrollTask?.disarm();
+    this.endScreenshotsOverlay();
   }
 
   /**

browser/components/screenshots/ScreenshotsUtils.sys.mjs

@@ -78,16 +78,20 @@ export class ScreenshotsComponentParent extends JSWindowActorParent {
         await ScreenshotsUtils.copyScreenshotFromRegion(region, browser);
         ScreenshotsUtils.exit(browser);
         break;
-      case "Screenshots:DownloadScreenshot":
+      case "Screenshots:DownloadScreenshot": {
         ScreenshotsUtils.closePanel(browser);
         ({ title, region } = message.data);
-        await ScreenshotsUtils.downloadScreenshotFromRegion(
-          title,
-          region,
-          browser
-        );
-        ScreenshotsUtils.exit(browser);
+        let downloadSucceeded =
+          await ScreenshotsUtils.downloadScreenshotFromRegion(
+            title,
+            region,
+            browser
+          );
+        if (downloadSucceeded) {
+          ScreenshotsUtils.exit(browser);
+        }
         break;
+      }
       case "Screenshots:OverlaySelection":
         ScreenshotsUtils.setPerBrowserState(browser, {
           hasOverlaySelection: message.data.hasSelection,
@@ -993,12 +997,6 @@ export var ScreenshotsUtils = {
    * @param type The type of screenshot taken.
    */
   async takeScreenshot(browser, type) {
-    this.closePanel(browser);
-    this.closeOverlay(browser, {
-      doNotResetMethods: true,
-      highlightRegions: true,
-    });
-
     Services.focus.setFocus(browser, 0);
 
     let rect;
@@ -1011,6 +1009,13 @@ export var ScreenshotsUtils = {
       lastUsedMethod = "visible";
     }
 
+    // Wait to close overlay until after we fetched bounds
+    this.closePanel(browser);
+    this.closeOverlay(browser, {
+      doNotResetMethods: true,
+      highlightRegions: true,
+    });
+
     let canvas = await this.createCanvas(rect, browser);
     let url = canvas.toDataURL();
 
@@ -1202,24 +1207,27 @@ export var ScreenshotsUtils = {
    * @param title The title of the current page
    * @param region The bounds of the screenshot
    * @param browser The current browser
+   * @returns {Promise<boolean>} true if the download succeeds, otherwise false
    */
   async downloadScreenshotFromRegion(title, region, browser) {
     let canvas = await this.createCanvas(region, browser);
     let dataUrl = canvas.toDataURL();
 
-    await this.downloadScreenshot(title, dataUrl, browser, {
+    return this.downloadScreenshot(title, dataUrl, browser, {
       object: "overlay_download",
     });
   },
 
   /**
-   * Download the screenshot
-   * This is called from the preview dialog
+   * Download the screenshot.
+   * This will return false if the download doesn't succeed for any reason and
+   * will return true if the screenshot is downloaded successfully.
+   *
    * @param title The title of the current page or null and getFilename will get the title
    * @param dataUrl The image data
    * @param browser The current browser
    * @param data Telemetry data
-   * @returns true if the download succeeds, otherwise false
+   * @returns {Promise<boolean>} true if the downnload succeeded, otherwise false
    */
   async downloadScreenshot(title, dataUrl, browser, data) {
     // Guard against missing image data.

browser/components/screenshots/tests/browser/browser_screenshots_test_downloads.js

@@ -240,3 +240,34 @@ add_task(async function test_download_filepicker_canceled() {
     }
   );
 });
+
+add_task(async function test_overlay_download_filepicker_canceled() {
+  await SpecialPowers.pushPrefEnv({
+    set: [["browser.download.useDownloadDir", false]],
+  });
+
+  await BrowserTestUtils.withNewTab(
+    {
+      gBrowser,
+      url: TEST_PAGE,
+    },
+    async browser => {
+      let helper = new ScreenshotsHelper(browser);
+
+      helper.triggerUIFromToolbar();
+      await helper.waitForOverlay();
+
+      await helper.dragOverlay(100, 100, 500, 500);
+
+      let filePickerCanceled = waitForFilePickerCancel();
+      await helper.clickDownloadButton();
+      info("download button clicked");
+      await filePickerCanceled;
+
+      ok(
+        await helper.isOverlayInitialized(),
+        "Overlay is still open and initialized"
+      );
+    }
+  );
+});