Bug 1477015 - Select storage.local backend on startup when the extension is not migrating its data. r=mixedpuppy,aswan draft
authorLuca Greco <lgreco@mozilla.com>
Thu, 26 Jul 2018 13:53:22 +0200
changeset 824696 64f70f7cdf9db304e60e501701e3b7f4161bf2d1
parent 822740 02c8644c45b1f143263d30d769d86c2d1058812e
push id117990
push userluca.greco@alcacoop.it
push dateTue, 31 Jul 2018 18:23:36 +0000
reviewersmixedpuppy, aswan
bugs1477015
milestone63.0a1
Bug 1477015 - Select storage.local backend on startup when the extension is not migrating its data. r=mixedpuppy,aswan MozReview-Commit-ID: WzW2bFlYNg
toolkit/components/extensions/Extension.jsm
toolkit/components/extensions/ExtensionStorageIDB.jsm
toolkit/components/extensions/child/ext-storage.js
toolkit/components/extensions/parent/ext-storage.js
toolkit/components/extensions/test/xpcshell/test_ext_storage_tab.js
toolkit/components/extensions/test/xpcshell/test_ext_storage_telemetry.js
--- a/toolkit/components/extensions/Extension.jsm
+++ b/toolkit/components/extensions/Extension.jsm
@@ -39,16 +39,17 @@ XPCOMUtils.defineLazyModuleGetters(this,
   AddonManager: "resource://gre/modules/AddonManager.jsm",
   AddonManagerPrivate: "resource://gre/modules/AddonManager.jsm",
   AddonSettings: "resource://gre/modules/addons/AddonSettings.jsm",
   AppConstants: "resource://gre/modules/AppConstants.jsm",
   AsyncShutdown: "resource://gre/modules/AsyncShutdown.jsm",
   ContextualIdentityService: "resource://gre/modules/ContextualIdentityService.jsm",
   ExtensionPermissions: "resource://gre/modules/ExtensionPermissions.jsm",
   ExtensionStorage: "resource://gre/modules/ExtensionStorage.jsm",
+  ExtensionStorageIDB: "resource://gre/modules/ExtensionStorageIDB.jsm",
   ExtensionTestCommon: "resource://testing-common/ExtensionTestCommon.jsm",
   FileSource: "resource://gre/modules/L10nRegistry.jsm",
   L10nRegistry: "resource://gre/modules/L10nRegistry.jsm",
   Log: "resource://gre/modules/Log.jsm",
   MessageChannel: "resource://gre/modules/MessageChannel.jsm",
   NetUtil: "resource://gre/modules/NetUtil.jsm",
   OS: "resource://gre/modules/osfile.jsm",
   PluralForm: "resource://gre/modules/PluralForm.jsm",
@@ -142,17 +143,16 @@ function classifyPermission(perm) {
   }
   return {permission: perm};
 }
 
 const LOGGER_ID_BASE = "addons.webextension.";
 const UUID_MAP_PREF = "extensions.webextensions.uuids";
 const LEAVE_STORAGE_PREF = "extensions.webextensions.keepStorageOnUninstall";
 const LEAVE_UUID_PREF = "extensions.webextensions.keepUuidOnUninstall";
-const IDB_MIGRATED_PREF_BRANCH = "extensions.webextensions.ExtensionStorageIDB.migrated";
 
 const COMMENT_REGEXP = new RegExp(String.raw`
     ^
     (
       (?:
         [^"\n] |
         " (?:[^"\\\n] | \\.)* "
       )*?
@@ -247,18 +247,17 @@ var UninstallObserver = {
       Services.qms.clearStoragesForPrincipal(principal);
 
       // Clear any storage.local data stored in the IDBBackend.
       let storagePrincipal = Services.scriptSecurityManager.createCodebasePrincipal(baseURI, {
         userContextId: WEBEXT_STORAGE_USER_CONTEXT_ID,
       });
       Services.qms.clearStoragesForPrincipal(storagePrincipal);
 
-      // Clear the preference set for the extensions migrated to the IDBBackend.
-      Services.prefs.clearUserPref(`${IDB_MIGRATED_PREF_BRANCH}.${addon.id}`);
+      ExtensionStorageIDB.clearMigratedExtensionPref(addon.id);
 
       // Clear localStorage created by the extension
       let storage = Services.domStorageManager.getStorage(null, principal);
       if (storage) {
         storage.clear();
       }
 
       // Remove any permissions related to the unlimitedStorage permission
@@ -1768,16 +1767,33 @@ class Extension extends ExtensionData {
       this.initSharedData();
 
       this.policy.active = false;
       this.policy = processScript.initExtension(this);
       this.policy.extension = this;
 
       this.updatePermissions(this.startupReason);
 
+      // Select the storage.local backend if it is already known,
+      // and start the data migration if needed.
+      if (this.hasPermission("storage")) {
+        if (!ExtensionStorageIDB.isBackendEnabled) {
+          this.setSharedData("storageIDBBackend", false);
+        } else if (ExtensionStorageIDB.isMigratedExtension(this)) {
+          this.setSharedData("storageIDBBackend", true);
+          this.setSharedData("storageIDBPrincipal", ExtensionStorageIDB.getStoragePrincipal(this));
+        } else {
+          // If the extension has to migrate backend, ensure that the data migration
+          // starts once Firefox is idle after the extension has been started.
+          this.once("ready", () => ChromeUtils.idleDispatch(() => {
+            ExtensionStorageIDB.selectBackend({extension: this});
+          }));
+        }
+      }
+
       // The "startup" Management event sent on the extension instance itself
       // is emitted just before the Management "startup" event,
       // and it is used to run code that needs to be executed before
       // any of the "startup" listeners.
       this.emit("startup", this);
       Management.emit("startup", this);
 
       await this.runManifest(this.manifest);
--- a/toolkit/components/extensions/ExtensionStorageIDB.jsm
+++ b/toolkit/components/extensions/ExtensionStorageIDB.jsm
@@ -368,31 +368,30 @@ async function migrateJSONFileData(exten
   let oldStorageExists;
   let idbConn;
   let jsonFile;
   let hasEmptyIDB;
   let nonFatalError;
   let dataMigrateCompleted = false;
   let hasOldData = false;
 
-  const isMigratedExtension = Services.prefs.getBoolPref(`${IDB_MIGRATED_PREF_BRANCH}.${extension.id}`, false);
-  if (isMigratedExtension) {
+  if (ExtensionStorageIDB.isMigratedExtension(extension)) {
     return;
   }
 
   try {
     idbConn = await ExtensionStorageLocalIDB.openForPrincipal(storagePrincipal);
     hasEmptyIDB = await idbConn.isEmpty();
 
     if (!hasEmptyIDB) {
       // If the IDB backend is enabled and there is data already stored in the IDB backend,
       // there is no "going back": any data that has not been migrated will be still on disk
       // but it is not going to be migrated anymore, it could be eventually used to allow
       // a user to manually retrieve the old data file).
-      Services.prefs.setBoolPref(`${IDB_MIGRATED_PREF_BRANCH}.${extension.id}`, true);
+      ExtensionStorageIDB.setMigratedExtensionPref(extension, true);
       return;
     }
   } catch (err) {
     extension.logWarning(
       `storage.local data migration cancelled, unable to open IDB connection: ${err.message}::${err.stack}`);
 
     DataMigrationTelemetry.recordResult({
       backend: "JSONFile",
@@ -480,17 +479,17 @@ async function migrateJSONFileData(exten
         await OS.File.move(oldStoragePath, openInfo.path);
       }
     } catch (err) {
       nonFatalError = err;
       extension.logWarning(err.message);
     }
   }
 
-  Services.prefs.setBoolPref(`${IDB_MIGRATED_PREF_BRANCH}.${extension.id}`, true);
+  ExtensionStorageIDB.setMigratedExtensionPref(extension, true);
 
   DataMigrationTelemetry.recordResult({
     backend: "IndexedDB",
     dataMigrated: dataMigrateCompleted,
     extensionId: extension.id,
     error: nonFatalError,
     hasJSONFile: oldStorageExists,
     hasOldData,
@@ -516,16 +515,28 @@ this.ExtensionStorageIDB = {
   //
   //   WeakMap<extension -> Promise<boolean>
   selectedBackendPromises: new WeakMap(),
 
   init() {
     XPCOMUtils.defineLazyPreferenceGetter(this, "isBackendEnabled", BACKEND_ENABLED_PREF, false);
   },
 
+  isMigratedExtension(extension) {
+    return Services.prefs.getBoolPref(`${IDB_MIGRATED_PREF_BRANCH}.${extension.id}`, false);
+  },
+
+  setMigratedExtensionPref(extension, val) {
+    Services.prefs.setBoolPref(`${IDB_MIGRATED_PREF_BRANCH}.${extension.id}`, !!val);
+  },
+
+  clearMigratedExtensionPref(extensionId) {
+    Services.prefs.clearUserPref(`${IDB_MIGRATED_PREF_BRANCH}.${extensionId}`);
+  },
+
   getStoragePrincipal(extension) {
     return extension.createPrincipal(extension.baseURI, {
       userContextId: WEBEXT_STORAGE_USER_CONTEXT_ID,
     });
   },
 
   /**
    * Select the preferred backend and return a promise which is resolved once the
--- a/toolkit/components/extensions/child/ext-storage.js
+++ b/toolkit/components/extensions/child/ext-storage.js
@@ -1,14 +1,16 @@
 "use strict";
 
 ChromeUtils.defineModuleGetter(this, "ExtensionStorage",
                                "resource://gre/modules/ExtensionStorage.jsm");
 ChromeUtils.defineModuleGetter(this, "ExtensionStorageIDB",
                                "resource://gre/modules/ExtensionStorageIDB.jsm");
+ChromeUtils.defineModuleGetter(this, "Services",
+                               "resource://gre/modules/Services.jsm");
 ChromeUtils.defineModuleGetter(this, "TelemetryStopwatch",
                                "resource://gre/modules/TelemetryStopwatch.jsm");
 
 // Telemetry histogram keys for the JSONFile backend.
 const storageGetHistogram = "WEBEXT_STORAGE_LOCAL_GET_MS";
 const storageSetHistogram = "WEBEXT_STORAGE_LOCAL_SET_MS";
 // Telemetry  histogram keys for the IndexedDB backend.
 const storageGetIDBHistogram = "WEBEXT_STORAGE_LOCAL_IDB_GET_MS";
@@ -50,17 +52,17 @@ this.storage = class extends ExtensionAP
       },
       clear() {
         return context.childManager.callParentAsyncFunction(
           "storage.local.JSONFileBackend.clear", []);
       },
     };
   }
 
-  getLocalIDBBackend(context, {hasParentListeners, serialize, storagePrincipal}) {
+  getLocalIDBBackend(context, {fireOnChanged, serialize, storagePrincipal}) {
     let dbPromise;
     async function getDB() {
       if (dbPromise) {
         return dbPromise;
       }
 
       dbPromise = ExtensionStorageIDB.open(storagePrincipal).catch(err => {
         // Reset the cached promise if it has been rejected, so that the next
@@ -81,59 +83,42 @@ this.storage = class extends ExtensionAP
       },
       set(items) {
         return measureOp(storageSetIDBHistogram, async () => {
           const db = await getDB();
           const changes = await db.set(items, {
             serialize: ExtensionStorage.serialize,
           });
 
-          if (!changes) {
-            return;
-          }
-
-          const hasListeners = await hasParentListeners();
-          if (hasListeners) {
-            await context.childManager.callParentAsyncFunction(
-              "storage.local.IDBBackend.fireOnChanged", [changes]);
+          if (changes) {
+            fireOnChanged(changes);
           }
         });
       },
       async remove(keys) {
         const db = await getDB();
         const changes = await db.remove(keys);
 
-        if (!changes) {
-          return;
-        }
-
-        const hasListeners = await hasParentListeners();
-        if (hasListeners) {
-          await context.childManager.callParentAsyncFunction(
-            "storage.local.IDBBackend.fireOnChanged", [changes]);
+        if (changes) {
+          fireOnChanged(changes);
         }
       },
       async clear() {
         const db = await getDB();
         const changes = await db.clear(context.extension);
 
-        if (!changes) {
-          return;
-        }
-
-        const hasListeners = await hasParentListeners();
-        if (hasListeners) {
-          await context.childManager.callParentAsyncFunction(
-            "storage.local.IDBBackend.fireOnChanged", [changes]);
+        if (changes) {
+          fireOnChanged(changes);
         }
       },
     };
   }
 
   getAPI(context) {
+    const {extension} = context;
     const serialize = ExtensionStorage.serializeForContext.bind(null, context);
     const deserialize = ExtensionStorage.deserializeForContext.bind(null, context);
 
     function sanitize(items) {
       // The schema validator already takes care of arrays (which are only allowed
       // to contain strings). Strings and null are safe values.
       if (typeof items != "object" || items === null || Array.isArray(items)) {
         return items;
@@ -147,58 +132,91 @@ this.storage = class extends ExtensionAP
       // So we enumerate all properties and sanitize each value individually.
       let sanitized = {};
       for (let [key, value] of Object.entries(items)) {
         sanitized[key] = ExtensionStorage.sanitize(value, context);
       }
       return sanitized;
     }
 
-    // Detect the actual storage.local enabled backend for the extension (as soon as the
-    // storage.local API has been accessed for the first time).
-    let promiseStorageLocalBackend;
+    function fireOnChanged(changes) {
+      // This call is used (by the storage.local API methods for the IndexedDB backend) to fire a storage.onChanged event,
+      // it uses the underlying message manager since the child context (or its ProxyContentParent counterpart
+      // running in the main process) may be gone by the time we call this, and so we can't use the childManager
+      // abstractions (e.g. callParentAsyncFunction or callParentFunctionNoReturn).
+      Services.cpmm.sendAsyncMessage(`Extension:StorageLocalOnChanged:${extension.uuid}`, changes);
+    }
+
+    // If the selected backend for the extension is not known yet, we have to lazily detect it
+    // by asking to the main process (as soon as the storage.local API has been accessed for
+    // the first time).
     const getStorageLocalBackend = async () => {
       const {
         backendEnabled,
         storagePrincipal,
       } = await ExtensionStorageIDB.selectBackend(context);
 
       if (!backendEnabled) {
         return this.getLocalFileBackend(context, {deserialize, serialize});
       }
 
       return this.getLocalIDBBackend(context, {
         storagePrincipal,
-        hasParentListeners() {
-          // We spare a good amount of memory if there are no listeners around
-          // (e.g. because they have never been subscribed or they have been removed
-          // in the meantime).
-          return context.childManager.callParentAsyncFunction(
-            "storage.local.IDBBackend.hasListeners", []);
-        },
+        fireOnChanged,
         serialize,
       });
     };
 
+    // Synchronously select the backend if it is already known.
+    let selectedBackend;
+
+    const useStorageIDBBackend = extension.getSharedData("storageIDBBackend");
+    if (useStorageIDBBackend === false) {
+      selectedBackend = this.getLocalFileBackend(context, {deserialize, serialize});
+    } else if (useStorageIDBBackend === true) {
+      selectedBackend = this.getLocalIDBBackend(context, {
+        storagePrincipal: extension.getSharedData("storageIDBPrincipal"),
+        fireOnChanged,
+        serialize,
+      });
+    }
+
+    let promiseStorageLocalBackend;
+
     // Generate the backend-agnostic local API wrapped methods.
     const local = {};
     for (let method of ["get", "set", "remove", "clear"]) {
       local[method] = async function(...args) {
         try {
-          if (!promiseStorageLocalBackend) {
-            promiseStorageLocalBackend = getStorageLocalBackend();
+          // Discover the selected backend if it is not known yet.
+          if (!selectedBackend) {
+            if (!promiseStorageLocalBackend) {
+              promiseStorageLocalBackend = getStorageLocalBackend().catch(err => {
+                // Clear the cached promise if it has been rejected.
+                promiseStorageLocalBackend = null;
+                throw err;
+              });
+            }
+
+            // If the storage.local method is not 'get' (which doesn't change any of the stored data),
+            // fall back to call the method in the parent process, so that it can be completed even
+            // if this context has been destroyed in the meantime.
+            if (method !== "get") {
+              // Let the outer try to catch rejections returned by the backend methods.
+              const result = await context.childManager.callParentAsyncFunction(
+                "storage.local.callMethodInParentProcess", [method, args]);
+              return result;
+            }
+
+            // Get the selected backend and cache it for the next API calls from this context.
+            selectedBackend = await promiseStorageLocalBackend;
           }
-          const backend = await promiseStorageLocalBackend.catch(err => {
-            // Clear the cached promise if it has been rejected.
-            promiseStorageLocalBackend = null;
-            throw err;
-          });
 
           // Let the outer try to catch rejections returned by the backend methods.
-          const result = await backend[method](...args);
+          const result = await selectedBackend[method](...args);
           return result;
         } catch (err) {
           // Ensure that the error we throw is converted into an ExtensionError
           // (e.g. DataCloneError instances raised from the internal IndexedDB
           // operation have to be converted to be accessible to the extension code).
           throw new ExtensionUtils.ExtensionError(String(err));
         }
       };
--- a/toolkit/components/extensions/parent/ext-storage.js
+++ b/toolkit/components/extensions/parent/ext-storage.js
@@ -29,22 +29,62 @@ const lookupManagedStorage = async (exte
   let info = await NativeManifests.lookupManifest("storage", extensionId, context);
   if (info) {
     return ExtensionStorage._serializableMap(info.manifest.data);
   }
   return null;
 };
 
 this.storage = class extends ExtensionAPI {
+  constructor(extension) {
+    super(extension);
+
+    const messageName = `Extension:StorageLocalOnChanged:${extension.uuid}`;
+    Services.ppmm.addMessageListener(messageName, this);
+    this.clearStorageChangedListener = () => {
+      Services.ppmm.removeMessageListener(messageName, this);
+    };
+  }
+
+  onShutdown() {
+    const {clearStorageChangedListener} = this;
+    this.clearStorageChangedListener = null;
+
+    if (clearStorageChangedListener) {
+      clearStorageChangedListener();
+    }
+  }
+
+  receiveMessage({name, data}) {
+    if (name !== `Extension:StorageLocalOnChanged:${this.extension.uuid}`) {
+      return;
+    }
+
+    ExtensionStorageIDB.notifyListeners(this.extension.id, data);
+  }
+
   getAPI(context) {
     let {extension} = context;
 
     return {
       storage: {
         local: {
+          async callMethodInParentProcess(method, args) {
+            const res = await ExtensionStorageIDB.selectBackend({extension});
+            if (!res.backendEnabled) {
+              return ExtensionStorage[method](extension.id, ...args);
+            }
+
+            const db = await ExtensionStorageIDB.open(res.storagePrincipal.deserialize(this));
+            const changes = await db[method](...args);
+            if (changes) {
+              ExtensionStorageIDB.notifyListeners(extension.id, changes);
+            }
+            return changes;
+          },
           // Private storage.local JSONFile backend methods (used internally by the child
           // ext-storage.js module).
           JSONFileBackend: {
             get(spec) {
               return ExtensionStorage.get(extension.id, spec);
             },
             set(items) {
               return ExtensionStorage.set(extension.id, items);
@@ -57,25 +97,16 @@ this.storage = class extends ExtensionAP
             },
           },
           // Private storage.local IDB backend methods (used internally by the child ext-storage.js
           // module).
           IDBBackend: {
             selectBackend() {
               return ExtensionStorageIDB.selectBackend(context);
             },
-            hasListeners() {
-              return ExtensionStorageIDB.hasListeners(extension.id);
-            },
-            fireOnChanged(changes) {
-              ExtensionStorageIDB.notifyListeners(extension.id, changes);
-            },
-            onceDataMigrated() {
-              return ExtensionStorageIDB.onceDataMigrated(context);
-            },
           },
         },
 
         sync: {
           get(spec) {
             enforceNoTemporaryAddon(extension.id);
             return extensionStorageSync.get(extension, spec, context);
           },
--- a/toolkit/components/extensions/test/xpcshell/test_ext_storage_tab.js
+++ b/toolkit/components/extensions/test/xpcshell/test_ext_storage_tab.js
@@ -99,28 +99,41 @@ add_task(async function test_storage_loc
 add_task(async function test_storage_local_idb_backend_from_tab() {
   return runWithPrefs([[ExtensionStorageIDB.BACKEND_ENABLED_PREF, true]],
                       test_multiple_pages);
 });
 
 async function test_storage_local_call_from_destroying_context() {
   let extension = ExtensionTestUtils.loadExtension({
     async background() {
+      let numberOfChanges = 0;
+      browser.storage.onChanged.addListener((changes, areaName) => {
+        if (areaName !== "local") {
+          browser.test.fail(`Received unexpected storage changes for "${areaName}"`);
+        }
+
+        numberOfChanges++;
+      });
+
       browser.test.onMessage.addListener(async ({msg, values}) => {
         switch (msg) {
           case "storage-set": {
             await browser.storage.local.set(values);
             browser.test.sendMessage("storage-set:done");
             break;
           }
           case "storage-get": {
             const res = await browser.storage.local.get();
             browser.test.sendMessage("storage-get:done", res);
             break;
           }
+          case "storage-changes": {
+            browser.test.sendMessage("storage-changes-count", numberOfChanges);
+            break;
+          }
           default:
             browser.test.fail(`Received unexpected message: ${msg}`);
         }
       });
 
       browser.test.sendMessage("ext-page-url", browser.runtime.getURL("tab.html"));
     },
     files: {
@@ -130,46 +143,52 @@ async function test_storage_local_call_f
             <meta charset="utf-8">
             <script src="tab.js"></script>
           </head>
         </html>`,
 
       "tab.js"() {
         browser.test.log("Extension tab - calling storage.local API method");
         // Call the storage.local API from a tab that is going to be quickly closed.
-        browser.storage.local.get({}).then(() => {
-          // This call should never be reached (because the tab should have been
-          // destroyed in the meantime).
-          browser.test.fail("Extension tab - Unexpected storage.local promise resolved");
+        browser.storage.local.set({
+          "test-key-from-destroying-context": "testvalue2",
         });
         // Navigate away from the extension page, so that the storage.local API call will be unable
         // to send the call to the caller context (because it has been destroyed in the meantime).
         window.location = "about:blank";
       },
     },
     manifest: {
       permissions: ["storage"],
     },
   });
 
   await extension.startup();
   const url = await extension.awaitMessage("ext-page-url");
 
   let contentPage = await ExtensionTestUtils.loadContentPage(url, {extension});
-  let expectedData = {"test-key": "test-value"};
+  let expectedBackgroundPageData = {"test-key-from-background-page": "test-value"};
+  let expectedTabData = {"test-key-from-destroying-context": "testvalue2"};
 
   info("Call storage.local.set from the background page and wait it to be completed");
-  extension.sendMessage({msg: "storage-set", values: expectedData});
+  extension.sendMessage({msg: "storage-set", values: expectedBackgroundPageData});
   await extension.awaitMessage("storage-set:done");
 
   info("Call storage.local.get from the background page and wait it to be completed");
   extension.sendMessage({msg: "storage-get"});
   let res = await extension.awaitMessage("storage-get:done");
 
-  Assert.deepEqual(res, expectedData, "Got the expected data set in the storage.local backend");
+  Assert.deepEqual(res, {
+    ...expectedBackgroundPageData,
+    ...expectedTabData,
+  }, "Got the expected data set in the storage.local backend");
+
+  extension.sendMessage({msg: "storage-changes"});
+  equal(await extension.awaitMessage("storage-changes-count"), 2,
+        "Got the expected number of storage.onChanged event received");
 
   contentPage.close();
 
   await extension.unload();
 }
 
 add_task(async function test_storage_local_file_backend_destroyed_context_promise() {
   return runWithPrefs([[ExtensionStorageIDB.BACKEND_ENABLED_PREF, false]],
--- a/toolkit/components/extensions/test/xpcshell/test_ext_storage_telemetry.js
+++ b/toolkit/components/extensions/test/xpcshell/test_ext_storage_telemetry.js
@@ -9,16 +9,19 @@ const HISTOGRAM_JSON_IDS = [
 ];
 
 const HISTOGRAM_IDB_IDS = [
   "WEBEXT_STORAGE_LOCAL_IDB_SET_MS", "WEBEXT_STORAGE_LOCAL_IDB_GET_MS",
 ];
 
 const HISTOGRAM_IDS = [].concat(HISTOGRAM_JSON_IDS, HISTOGRAM_IDB_IDS);
 
+const EXTENSION_ID1 = "@test-extension1";
+const EXTENSION_ID2 = "@test-extension2";
+
 async function test_telemetry_background() {
   const expectedEmptyHistograms = ExtensionStorageIDB.isBackendEnabled ?
           HISTOGRAM_JSON_IDS : HISTOGRAM_IDB_IDS;
 
   const expectedNonEmptyHistograms = ExtensionStorageIDB.isBackendEnabled ?
           HISTOGRAM_IDB_IDS : HISTOGRAM_JSON_IDS;
 
   const server = createHttpServer();
@@ -27,42 +30,62 @@ async function test_telemetry_background
   const BASE_URL = `http://localhost:${server.identity.primaryPort}/data`;
 
   async function contentScript() {
     await browser.storage.local.set({a: "b"});
     await browser.storage.local.get("a");
     browser.runtime.sendMessage("contentDone");
   }
 
-  let extInfo = {
-    manifest: {
-      permissions: ["storage"],
-      content_scripts: [
-        {
-          "matches": ["http://*/*/file_sample.html"],
-          "js": ["content_script.js"],
-        },
-      ],
-    },
+  let baseManifest = {
+    permissions: ["storage"],
+    content_scripts: [
+      {
+        "matches": ["http://*/*/file_sample.html"],
+        "js": ["content_script.js"],
+      },
+    ],
+  };
+
+  let baseExtInfo = {
     async background() {
       browser.runtime.onMessage.addListener(msg => {
         browser.test.sendMessage(msg);
       });
 
       await browser.storage.local.set({a: "b"});
       await browser.storage.local.get("a");
       browser.test.sendMessage("backgroundDone");
     },
     files: {
       "content_script.js": contentScript,
     },
   };
 
-  let extension1 = ExtensionTestUtils.loadExtension(extInfo);
-  let extension2 = ExtensionTestUtils.loadExtension(extInfo);
+  let extInfo1 = {
+    ...baseExtInfo,
+    manifest: {
+      ...baseManifest,
+      applications: {
+        gecko: {id: EXTENSION_ID1},
+      },
+    },
+  };
+  let extInfo2 = {
+    ...baseExtInfo,
+    manifest: {
+      ...baseManifest,
+      applications: {
+        gecko: {id: EXTENSION_ID2},
+      },
+    },
+  };
+
+  let extension1 = ExtensionTestUtils.loadExtension(extInfo1);
+  let extension2 = ExtensionTestUtils.loadExtension(extInfo2);
 
   clearHistograms();
 
   let process = IS_OOP ? "extension" : "parent";
   let snapshots = getSnapshots(process);
 
   for (let id of HISTOGRAM_IDS) {
     ok(!(id in snapshots), `No data recorded for histogram: ${id}.`);
@@ -124,11 +147,18 @@ async function test_telemetry_background
 }
 
 add_task(function test_telemetry_background_file_backend() {
   return runWithPrefs([[ExtensionStorageIDB.BACKEND_ENABLED_PREF, false]],
                       test_telemetry_background);
 });
 
 add_task(function test_telemetry_background_idb_backend() {
-  return runWithPrefs([[ExtensionStorageIDB.BACKEND_ENABLED_PREF, true]],
-                      test_telemetry_background);
+  return runWithPrefs([
+    [ExtensionStorageIDB.BACKEND_ENABLED_PREF, true],
+    // Set the migrated preference for the two test extension, because the
+    // first storage.local call fallbacks to run in the parent process when we
+    // don't know which is the selected backend during the extension startup
+    // and so we can't choose the telemetry histogram to use.
+    [`${ExtensionStorageIDB.IDB_MIGRATED_PREF_BRANCH}.${EXTENSION_ID1}`, true],
+    [`${ExtensionStorageIDB.IDB_MIGRATED_PREF_BRANCH}.${EXTENSION_ID2}`, true],
+  ], test_telemetry_background);
 });