Bug 1722567 - Save group of closed tabs to restore the all group. r=kashav

When a group of tabs is closed, save the it in session data so tabs could be restored together.

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

browser/base/content/browser.js

@@ -8273,7 +8273,6 @@ function undoCloseTab(aIndex) {
       }
     }
   }
-  SessionStore.setLastClosedTabCount(window, 1);
 
   return tab;
 }

browser/base/content/tabbrowser.js

@@ -2583,7 +2583,6 @@
       document
         .getElementById("History:UndoCloseTab")
         .setAttribute("data-l10n-args", JSON.stringify({ tabCount: 1 }));
-      SessionStore.setLastClosedTabCount(window, 1);
 
       // if we're adding tabs, we're past interrupt mode, ditch the owner
       if (this.selectedTab.owner) {
@@ -3425,7 +3424,7 @@
         return;
       }
 
-      let initialTabCount = tabs.length;
+      SessionStore.resetLastClosedTabCount(window);
       this._clearMultiSelectionLocked = true;
 
       // Guarantee that _clearMultiSelectionLocked lock gets released.
@@ -3437,6 +3436,7 @@
         let aParams = { animate, prewarmed: true };
 
         for (let tab of tabs) {
+          tab._closedInGroup = true;
           if (tab.selected) {
             lastToClose = tab;
             let toBlurTo = this._findTabToBlurTo(lastToClose, tabs);
@@ -3516,6 +3516,10 @@
         // Now run again sequentially the beforeunload listeners that will result in a prompt.
         for (let tab of tabsWithBeforeUnloadPrompt) {
           this.removeTab(tab, aParams);
+          if (!tab.closing) {
+            // If we abort the closing of the tab.
+            tab._closedInGroup = false;
+          }
         }
 
         // Avoid changing the selected browser several times by removing it,
@@ -3529,17 +3533,14 @@
 
       this._clearMultiSelectionLocked = false;
       this.avoidSingleSelectedTab();
-      let closedTabsCount =
-        initialTabCount - tabs.filter(t => t.isConnected && !t.closing).length;
       // Don't use document.l10n.setAttributes because the FTL file is loaded
       // lazily and we won't be able to resolve the string.
-      document
+      document.getElementById("History:UndoCloseTab").setAttribute(
-        .getElementById("History:UndoCloseTab")
+        "data-l10n-args",
-        .setAttribute(
+        JSON.stringify({
-          "data-l10n-args",
+          tabCount: SessionStore.getLastClosedTabCount(window),
-          JSON.stringify({ tabCount: closedTabsCount })
+        })
-        );
+      );
-      SessionStore.setLastClosedTabCount(window, closedTabsCount);
     },
 
     removeCurrentTab(aParams) {

browser/base/content/test/tabs/browser_undo_close_tabs.js

@@ -26,12 +26,12 @@ add_task(async function withMultiSelectedTabs() {
   ok(tab2.multiselected, "Tab2 is multiselected");
   ok(tab3.multiselected, "Tab3 is multiselected");
   ok(tab4.multiselected, "Tab4 is multiselected");
-  is(gBrowser.multiSelectedTabsCount, 3, "Two multiselected tabs");
+  is(gBrowser.multiSelectedTabsCount, 3, "Three multiselected tabs");
 
   gBrowser.removeMultiSelectedTabs();
   await TestUtils.waitForCondition(
-    () => gBrowser.tabs.length == 2,
+    () => SessionStore.getLastClosedTabCount(window) == 3,
-    "wait for the multiselected tabs to close"
+    "wait for the multi selected tabs to close in SessionStore"
   );
   is(
     SessionStore.getLastClosedTabCount(window),
@@ -44,6 +44,13 @@ add_task(async function withMultiSelectedTabs() {
     () => gBrowser.tabs.length == 5,
     "wait for the tabs to reopen"
   );
+
+  is(
+    SessionStore.getLastClosedTabCount(window),
+    SessionStore.getClosedTabCount(window) ? 1 : 0,
+    "LastClosedTabCount should be reset"
+  );
+
   info("waiting for the browsers to finish loading");
   // Check that the tabs are restored in the correct order
   for (let tabId of [2, 3, 4]) {
@@ -61,6 +68,67 @@ add_task(async function withMultiSelectedTabs() {
   gBrowser.removeAllTabsBut(initialTab);
 });
 
+add_task(async function withBothGroupsAndTab() {
+  let initialTab = gBrowser.selectedTab;
+  let tab1 = await addTab("https://example.com/1");
+  let tab2 = await addTab("https://example.com/2");
+  let tab3 = await addTab("https://example.com/3");
+
+  gBrowser.selectedTab = tab2;
+  await triggerClickOn(tab3, { shiftKey: true });
+
+  ok(!initialTab.multiselected, "InitialTab is not multiselected");
+  ok(!tab1.multiselected, "Tab1 is not multiselected");
+  ok(tab2.multiselected, "Tab2 is multiselected");
+  ok(tab3.multiselected, "Tab3 is multiselected");
+  is(gBrowser.multiSelectedTabsCount, 2, "Two multiselected tabs");
+
+  gBrowser.removeMultiSelectedTabs();
+  await TestUtils.waitForCondition(
+    () => gBrowser.tabs.length == 2,
+    "wait for the multiselected tabs to close"
+  );
+
+  is(
+    SessionStore.getLastClosedTabCount(window),
+    2,
+    "SessionStore should know how many tabs were just closed"
+  );
+
+  let tab4 = await addTab("http://example.com/4");
+
+  is(
+    SessionStore.getLastClosedTabCount(window),
+    2,
+    "LastClosedTabCount should be the same"
+  );
+
+  gBrowser.removeTab(tab4);
+
+  await TestUtils.waitForCondition(
+    () => SessionStore.getLastClosedTabCount(window) == 1,
+    "wait for the tab to close in SessionStore"
+  );
+
+  let count = 3;
+  for (let i = 0; i < 3; i++) {
+    is(
+      SessionStore.getLastClosedTabCount(window),
+      1,
+      "LastClosedTabCount should be one"
+    );
+    undoCloseTab();
+
+    await TestUtils.waitForCondition(
+      () => gBrowser.tabs.length == count,
+      "wait for the tabs to reopen"
+    );
+    count++;
+  }
+
+  gBrowser.removeAllTabsBut(initialTab);
+});
+
 add_task(async function withCloseTabsToTheRight() {
   let initialTab = gBrowser.selectedTab;
   let tab1 = await addTab("https://example.com/1");

browser/components/sessionstore/SessionStore.jsm

@@ -309,8 +309,9 @@ var SessionStore = {
   getLastClosedTabCount(aWindow) {
     return SessionStoreInternal.getLastClosedTabCount(aWindow);
   },
-  setLastClosedTabCount(aWindow, aNumber) {
+
-    return SessionStoreInternal.setLastClosedTabCount(aWindow, aNumber);
+  resetLastClosedTabCount(aWindow) {
+    SessionStoreInternal.resetLastClosedTabCount(aWindow);
   },
 
   getClosedTabCount: function ss_getClosedTabCount(aWindow) {
@@ -763,7 +764,6 @@ var SessionStoreInternal = {
 
     this._initPrefs();
     this._initialized = true;
-    this._closedTabCache = new WeakMap();
 
     Services.telemetry
       .getHistogramById("FX_SESSION_RESTORE_PRIVACY_LEVEL")
@@ -1198,7 +1198,7 @@ var SessionStoreInternal = {
       this._closedTabs.has(permanentKey) &&
       !this._crashedBrowsers.has(permanentKey)
     ) {
-      let { closedTabs, tabData } = this._closedTabs.get(permanentKey);
+      let { winData, closedTabs, tabData } = this._closedTabs.get(permanentKey);
 
       // We expect no further updates.
       this._closedTabs.delete(permanentKey);
@@ -1215,12 +1215,12 @@ var SessionStoreInternal = {
         // the list of closed tabs when it was closed (because we deemed
         // the state not worth saving) then add it to the window's list
         // of closed tabs now.
-        this.saveClosedTabData(closedTabs, tabData);
+        this.saveClosedTabData(winData, closedTabs, tabData);
       } else if (!shouldSave && index > -1) {
         // Remove from the list of closed tabs. The update messages sent
         // after the tab was closed changed enough state so that we no
         // longer consider its data interesting enough to keep around.
-        this.removeClosedTabData(closedTabs, index);
+        this.removeClosedTabData(winData, closedTabs, index);
       }
     }
 
@@ -1477,6 +1477,7 @@ var SessionStoreInternal = {
       tabs: [],
       selected: 0,
       _closedTabs: [],
+      _lastClosedTabGroupCount: -1,
       busy: false,
     };
 
@@ -2494,9 +2495,11 @@ var SessionStoreInternal = {
       image: aWindow.gBrowser.getIcon(aTab),
       pos: aTab._tPos,
       closedAt: Date.now(),
+      closedInGroup: aTab._closedInGroup,
     };
 
-    let closedTabs = this._windows[aWindow.__SSi]._closedTabs;
+    let winData = this._windows[aWindow.__SSi];
+    let closedTabs = winData._closedTabs;
 
     // Determine whether the tab contains any information worth saving. Note
     // that there might be pending state changes queued in the child that
@@ -2507,12 +2510,12 @@ var SessionStoreInternal = {
       // of the list but those cases should be extremely rare and
       // do probably never occur when using the browser normally.
       // (Tests or add-ons might do weird things though.)
-      this.saveClosedTabData(closedTabs, tabData);
+      this.saveClosedTabData(winData, closedTabs, tabData);
     }
 
     // Remember the closed tab to properly handle any last updates included in
     // the final "update" message sent by the frame script's unload handler.
-    this._closedTabs.set(permanentKey, { closedTabs, tabData });
+    this._closedTabs.set(permanentKey, { winData, closedTabs, tabData });
   },
 
   /**
@@ -2626,12 +2629,14 @@ var SessionStoreInternal = {
    * all tabs already in the list. The list will be truncated to contain a
    * maximum of |this._max_tabs_undo| entries.
    *
-   * @param closedTabs (array)
+   * @param winData (object)
-   *        The list of closed tabs for a window.
+   *        The data of the window.
    * @param tabData (object)
    *        The tabData to be inserted.
+   * @param closedTabs (array)
+   *        The list of closed tabs for a window.
    */
-  saveClosedTabData(closedTabs, tabData) {
+  saveClosedTabData(winData, closedTabs, tabData) {
     // Find the index of the first tab in the list
     // of closed tabs that was closed before our tab.
     let index = closedTabs.findIndex(tab => {
@@ -2651,6 +2656,18 @@ var SessionStoreInternal = {
     closedTabs.splice(index, 0, tabData);
     this._closedObjectsChanged = true;
 
+    if (tabData.closedInGroup) {
+      if (winData._lastClosedTabGroupCount < this._max_tabs_undo) {
+        if (winData._lastClosedTabGroupCount < 0) {
+          winData._lastClosedTabGroupCount = 1;
+        } else {
+          winData._lastClosedTabGroupCount++;
+        }
+      }
+    } else {
+      winData._lastClosedTabGroupCount = -1;
+    }
+
     // Truncate the list of closed tabs, if needed.
     if (closedTabs.length > this._max_tabs_undo) {
       closedTabs.splice(this._max_tabs_undo, closedTabs.length);
@@ -2662,16 +2679,24 @@ var SessionStoreInternal = {
    * the tab's final message is still pending we will simply discard it when
    * it arrives so that the tab doesn't reappear in the list.
    *
-   * @param closedTabs (array)
+   * @param winData (object)
-   *        The list of closed tabs for a window.
+   *        The data of the window.
    * @param index (uint)
    *        The index of the tab to remove.
+   * @param closedTabs (array)
+   *        The list of closed tabs for a window.
    */
-  removeClosedTabData(closedTabs, index) {
+  removeClosedTabData(winData, closedTabs, index) {
     // Remove the given index from the list.
     let [closedTab] = closedTabs.splice(index, 1);
     this._closedObjectsChanged = true;
 
+    // If the tab is part of the last closed group,
+    // we need to deduct the tab from the count.
+    if (index < winData._lastClosedTabGroupCount) {
+      winData._lastClosedTabGroupCount--;
+    }
+
     // If the closed tab's state still has a .permanentKey property then we
     // haven't seen its final update message yet. Remove it from the map of
     // closed tabs so that we will simply discard its last messages and will
@@ -3100,24 +3125,21 @@ var SessionStoreInternal = {
 
   getLastClosedTabCount(aWindow) {
     if ("__SSi" in aWindow) {
-      // Blank tabs cannot be undo-closed, so the number returned by
-      // the ClosedTabCache can be greater than the return value of
-      // getClosedTabCount. We won't restore blank tabs, so we return
-      // the minimum of these two values.
       return Math.min(
-        this._closedTabCache.get(aWindow) || 1,
+        Math.max(this._windows[aWindow.__SSi]._lastClosedTabGroupCount, 1),
         this.getClosedTabCount(aWindow)
       );
     }
 
     throw (Components.returnCode = Cr.NS_ERROR_INVALID_ARG);
   },
-  setLastClosedTabCount(aWindow, aNumber) {
-    if ("__SSi" in aWindow) {
-      return this._closedTabCache.set(aWindow, aNumber);
-    }
 
-    throw (Components.returnCode = Cr.NS_ERROR_INVALID_ARG);
+  resetLastClosedTabCount(aWindow) {
+    if ("__SSi" in aWindow) {
+      this._windows[aWindow.__SSi]._lastClosedTabGroupCount = -1;
+    } else {
+      throw (Components.returnCode = Cr.NS_ERROR_INVALID_ARG);
+    }
   },
 
   getClosedTabCount: function ssi_getClosedTabCount(aWindow) {
@@ -3163,11 +3185,11 @@ var SessionStoreInternal = {
       );
     }
 
-    var closedTabs = this._windows[aWindow.__SSi]._closedTabs;
+    let winData = this._windows[aWindow.__SSi];
 
     // default to the most-recently closed tab
     aIndex = aIndex || 0;
-    if (!(aIndex in closedTabs)) {
+    if (!(aIndex in winData._closedTabs)) {
       throw Components.Exception(
         "Invalid index: not in the closed tabs",
         Cr.NS_ERROR_INVALID_ARG
@@ -3175,7 +3197,11 @@ var SessionStoreInternal = {
     }
 
     // fetch the data of closed tab, while removing it from the array
-    let { state, pos } = this.removeClosedTabData(closedTabs, aIndex);
+    let { state, pos } = this.removeClosedTabData(
+      winData,
+      winData._closedTabs,
+      aIndex
+    );
 
     // create a new tab
     let tabbrowser = aWindow.gBrowser;
@@ -3202,11 +3228,11 @@ var SessionStoreInternal = {
       );
     }
 
-    var closedTabs = this._windows[aWindow.__SSi]._closedTabs;
+    let winData = this._windows[aWindow.__SSi];
 
     // default to the most-recently closed tab
     aIndex = aIndex || 0;
-    if (!(aIndex in closedTabs)) {
+    if (!(aIndex in winData._closedTabs)) {
       throw Components.Exception(
         "Invalid index: not in the closed tabs",
         Cr.NS_ERROR_INVALID_ARG
@@ -3214,7 +3240,7 @@ var SessionStoreInternal = {
     }
 
     // remove closed tab from the array
-    this.removeClosedTabData(closedTabs, aIndex);
+    this.removeClosedTabData(winData, winData._closedTabs, aIndex);
 
     // Notify of changes to closed objects.
     this._notifyOfClosedObjectsChange();
@@ -3499,7 +3525,7 @@ var SessionStoreInternal = {
         });
 
         for (let index of indexes.reverse()) {
-          this.removeClosedTabData(windowState._closedTabs, index);
+          this.removeClosedTabData(windowState, windowState._closedTabs, index);
         }
       }
     }