Bug 1350522: Part 6 - Cleanup per-api-instance state logic. r?aswan draft
authorKris Maglione <maglione.k@gmail.com>
Sat, 25 Mar 2017 11:36:56 -0700
changeset 551554 4dd4e9b3a12ff695f5a10f6674d99d56331eec10
parent 551553 3f146583e90c815198dfe522d4a9a7b76bf5d8db
child 551555 a6ee189146ffb3a4f011f16a6432dc228569a1a1
push id51079
push usermaglione.k@gmail.com
push dateMon, 27 Mar 2017 01:43:58 +0000
reviewersaswan
bugs1350522
milestone55.0a1
Bug 1350522: Part 6 - Cleanup per-api-instance state logic. r?aswan MozReview-Commit-ID: 5ixBA34fvFf
browser/components/extensions/ext-browserAction.js
browser/components/extensions/ext-commands.js
browser/components/extensions/ext-omnibox.js
browser/components/extensions/ext-pageAction.js
browser/components/extensions/ext-sidebarAction.js
toolkit/components/extensions/ext-backgroundPage.js
toolkit/components/extensions/ext-geolocation.js
toolkit/components/extensions/ext-protocolHandlers.js
toolkit/components/extensions/ext-theme.js
--- a/browser/components/extensions/ext-browserAction.js
+++ b/browser/components/extensions/ext-browserAction.js
@@ -462,133 +462,132 @@ BrowserAction.for = (extension) => {
 
 global.browserActionFor = BrowserAction.for;
 
 this.browserAction = class extends ExtensionAPI {
   onManifestEntry(entryName) {
     let {extension} = this;
     let {manifest} = extension;
 
-    let browserAction = new BrowserAction(manifest.browser_action, extension);
-    browserAction.build();
-    browserActionMap.set(extension, browserAction);
+    this.browserAction = new BrowserAction(manifest.browser_action, extension);
+    this.browserAction.build();
+    browserActionMap.set(extension, this.browserAction);
   }
 
   onShutdown(reason) {
     let {extension} = this;
 
-    if (browserActionMap.has(extension)) {
-      browserActionMap.get(extension).shutdown();
-      browserActionMap.delete(extension);
-    }
+    browserActionMap.delete(extension);
+    this.browserAction.shutdown();
   }
 
   getAPI(context) {
     let {extension} = context;
+    let {tabManager} = extension;
 
-    let {tabManager} = extension;
+    let {browserAction} = this;
 
     function getTab(tabId) {
       if (tabId !== null) {
         return tabTracker.getTab(tabId);
       }
       return null;
     }
 
     return {
       browserAction: {
         onClicked: new SingletonEventManager(context, "browserAction.onClicked", fire => {
           let listener = () => {
             fire.async(tabManager.convert(tabTracker.activeTab));
           };
-          BrowserAction.for(extension).on("click", listener);
+          browserAction.on("click", listener);
           return () => {
-            BrowserAction.for(extension).off("click", listener);
+            browserAction.off("click", listener);
           };
         }).api(),
 
         enable: function(tabId) {
           let tab = getTab(tabId);
-          BrowserAction.for(extension).setProperty(tab, "enabled", true);
+          browserAction.setProperty(tab, "enabled", true);
         },
 
         disable: function(tabId) {
           let tab = getTab(tabId);
-          BrowserAction.for(extension).setProperty(tab, "enabled", false);
+          browserAction.setProperty(tab, "enabled", false);
         },
 
         setTitle: function(details) {
           let tab = getTab(details.tabId);
 
           let title = details.title;
           // Clear the tab-specific title when given a null string.
           if (tab && title == "") {
             title = null;
           }
-          BrowserAction.for(extension).setProperty(tab, "title", title);
+          browserAction.setProperty(tab, "title", title);
         },
 
         getTitle: function(details) {
           let tab = getTab(details.tabId);
 
-          let title = BrowserAction.for(extension).getProperty(tab, "title");
+          let title = browserAction.getProperty(tab, "title");
           return Promise.resolve(title);
         },
 
         setIcon: function(details) {
           let tab = getTab(details.tabId);
 
           let icon = IconDetails.normalize(details, extension, context);
-          BrowserAction.for(extension).setProperty(tab, "icon", icon);
+          browserAction.setProperty(tab, "icon", icon);
         },
 
         setBadgeText: function(details) {
           let tab = getTab(details.tabId);
 
-          BrowserAction.for(extension).setProperty(tab, "badgeText", details.text);
+          browserAction.setProperty(tab, "badgeText", details.text);
         },
 
         getBadgeText: function(details) {
           let tab = getTab(details.tabId);
 
-          let text = BrowserAction.for(extension).getProperty(tab, "badgeText");
+          let text = browserAction.getProperty(tab, "badgeText");
           return Promise.resolve(text);
         },
 
         setPopup: function(details) {
           let tab = getTab(details.tabId);
 
           // Note: Chrome resolves arguments to setIcon relative to the calling
           // context, but resolves arguments to setPopup relative to the extension
           // root.
           // For internal consistency, we currently resolve both relative to the
           // calling context.
           let url = details.popup && context.uri.resolve(details.popup);
-          BrowserAction.for(extension).setProperty(tab, "popup", url);
+          browserAction.setProperty(tab, "popup", url);
         },
 
         getPopup: function(details) {
           let tab = getTab(details.tabId);
 
-          let popup = BrowserAction.for(extension).getProperty(tab, "popup");
+          let popup = browserAction.getProperty(tab, "popup");
           return Promise.resolve(popup);
         },
 
         setBadgeBackgroundColor: function(details) {
           let tab = getTab(details.tabId);
           let color = details.color;
           if (!Array.isArray(color)) {
             let col = DOMUtils.colorToRGBA(color);
             color = col && [col.r, col.g, col.b, Math.round(col.a * 255)];
           }
-          BrowserAction.for(extension).setProperty(tab, "badgeBackgroundColor", color);
+          browserAction.setProperty(tab, "badgeBackgroundColor", color);
         },
 
         getBadgeBackgroundColor: function(details, callback) {
           let tab = getTab(details.tabId);
 
-          let color = BrowserAction.for(extension).getProperty(tab, "badgeBackgroundColor");
+          let color = browserAction.getProperty(tab, "badgeBackgroundColor");
           return Promise.resolve(color || [0xd9, 0, 0, 255]);
         },
       },
     };
   }
 };
--- a/browser/components/extensions/ext-commands.js
+++ b/browser/components/extensions/ext-commands.js
@@ -4,19 +4,16 @@
 
 var {
   SingletonEventManager,
   PlatformInfo,
 } = ExtensionUtils;
 
 var XUL_NS = "http://www.mozilla.org/keymaster/gatekeeper/there.is.only.xul";
 
-// WeakMap[Extension -> CommandList]
-let commandsMap = new WeakMap();
-
 function CommandList(manifest, extension) {
   this.extension = extension;
   this.id = makeWidgetId(extension.id);
   this.windowOpenListener = null;
 
   // Map[{String} commandName -> {Object} commandProperties]
   this.commands = this.loadCommandsFromManifest(manifest);
 
@@ -220,48 +217,41 @@ CommandList.prototype = {
   },
 };
 
 this.commands = class extends ExtensionAPI {
   onManifestEntry(entryName) {
     let {extension} = this;
     let {manifest} = extension;
 
-    commandsMap.set(extension, new CommandList(manifest, extension));
+    this.commandList = new CommandList(manifest, extension);
   }
 
   onShutdown(reason) {
-    let {extension} = this;
-
-    let commandsList = commandsMap.get(extension);
-    if (commandsList) {
-      commandsList.unregister();
-      commandsMap.delete(extension);
-    }
+    this.commandList.unregister();
   }
 
   getAPI(context) {
-    let {extension} = context;
     return {
       commands: {
-        getAll() {
-          let commands = commandsMap.get(extension).commands;
+        getAll: () => {
+          let commands = this.commandList.commands;
           return Promise.resolve(Array.from(commands, ([name, command]) => {
             return ({
               name,
               description: command.description,
               shortcut: command.shortcut,
             });
           }));
         },
         onCommand: new SingletonEventManager(context, "commands.onCommand", fire => {
           let listener = (eventName, commandName) => {
             fire.async(commandName);
           };
-          commandsMap.get(extension).on("command", listener);
+          this.commandList.on("command", listener);
           return () => {
-            commandsMap.get(extension).off("command", listener);
+            this.commandList.off("command", listener);
           };
         }).api(),
       },
     };
   }
 };
--- a/browser/components/extensions/ext-omnibox.js
+++ b/browser/components/extensions/ext-omnibox.js
@@ -3,53 +3,43 @@
 "use strict";
 
 XPCOMUtils.defineLazyModuleGetter(this, "ExtensionSearchHandler",
                                   "resource://gre/modules/ExtensionSearchHandler.jsm");
 var {
   SingletonEventManager,
 } = ExtensionUtils;
 
-// WeakMap[extension -> keyword]
-let gKeywordMap = new WeakMap();
-
 this.omnibox = class extends ExtensionAPI {
   onManifestEntry(entryName) {
     let {extension} = this;
     let {manifest} = extension;
 
     let keyword = manifest.omnibox.keyword;
     try {
       // This will throw if the keyword is already registered.
       ExtensionSearchHandler.registerKeyword(keyword, extension);
-      gKeywordMap.set(extension, keyword);
+      this.keyword = keyword;
     } catch (e) {
       extension.manifestError(e.message);
     }
   }
 
   onShutdown(reason) {
-    let {extension} = this;
-
-    let keyword = gKeywordMap.get(extension);
-    if (keyword) {
-      ExtensionSearchHandler.unregisterKeyword(keyword);
-      gKeywordMap.delete(extension);
-    }
+    ExtensionSearchHandler.unregisterKeyword(this.keyword);
   }
 
   getAPI(context) {
     let {extension} = context;
     return {
       omnibox: {
-        setDefaultSuggestion(suggestion) {
-          let keyword = gKeywordMap.get(extension);
+        setDefaultSuggestion: (suggestion) => {
           try {
             // This will throw if the keyword failed to register.
-            ExtensionSearchHandler.setDefaultSuggestion(keyword, suggestion);
+            ExtensionSearchHandler.setDefaultSuggestion(this.keyword, suggestion);
           } catch (e) {
             return Promise.reject(e.message);
           }
         },
 
         onInputStarted: new SingletonEventManager(context, "omnibox.onInputStarted", fire => {
           let listener = (eventName) => {
             fire.sync();
@@ -77,20 +67,19 @@ this.omnibox = class extends ExtensionAP
           extension.on(ExtensionSearchHandler.MSG_INPUT_ENTERED, listener);
           return () => {
             extension.off(ExtensionSearchHandler.MSG_INPUT_ENTERED, listener);
           };
         }).api(),
       },
 
       omnibox_internal: {
-        addSuggestions(id, suggestions) {
-          let keyword = gKeywordMap.get(extension);
+        addSuggestions: (id, suggestions) => {
           try {
-            ExtensionSearchHandler.addSuggestions(keyword, id, suggestions);
+            ExtensionSearchHandler.addSuggestions(this.keyword, id, suggestions);
           } catch (e) {
             // Silently fail because the extension developer can not know for sure if the user
             // has already invalidated the callback when asynchronously providing suggestions.
           }
         },
 
         onInputChanged: new SingletonEventManager(context, "omnibox_internal.onInputChanged", fire => {
           let listener = (eventName, text, id) => {
--- a/browser/components/extensions/ext-pageAction.js
+++ b/browser/components/extensions/ext-pageAction.js
@@ -229,93 +229,89 @@ PageAction.for = extension => {
 
 global.pageActionFor = PageAction.for;
 
 this.pageAction = class extends ExtensionAPI {
   onManifestEntry(entryName) {
     let {extension} = this;
     let {manifest} = extension;
 
-    let pageAction = new PageAction(manifest.page_action, extension);
-    pageActionMap.set(extension, pageAction);
+    this.pageAction = new PageAction(manifest.page_action, extension);
+    pageActionMap.set(extension, this.pageAction);
   }
 
   onShutdown(reason) {
-    let {extension} = this;
-
-    if (pageActionMap.has(extension)) {
-      pageActionMap.get(extension).shutdown();
-      pageActionMap.delete(extension);
-    }
+    pageActionMap.delete(this.extension);
+    this.pageAction.shutdown();
   }
 
   getAPI(context) {
     let {extension} = context;
 
     const {tabManager} = extension;
+    const {pageAction} = this;
 
     return {
       pageAction: {
         onClicked: new SingletonEventManager(context, "pageAction.onClicked", fire => {
           let listener = (evt, tab) => {
             fire.async(tabManager.convert(tab));
           };
-          let pageAction = PageAction.for(extension);
 
           pageAction.on("click", listener);
           return () => {
             pageAction.off("click", listener);
           };
         }).api(),
 
         show(tabId) {
           let tab = tabTracker.getTab(tabId);
-          PageAction.for(extension).setProperty(tab, "show", true);
+          pageAction.setProperty(tab, "show", true);
         },
 
         hide(tabId) {
           let tab = tabTracker.getTab(tabId);
-          PageAction.for(extension).setProperty(tab, "show", false);
+          pageAction.setProperty(tab, "show", false);
         },
 
         setTitle(details) {
           let tab = tabTracker.getTab(details.tabId);
 
           // Clear the tab-specific title when given a null string.
-          PageAction.for(extension).setProperty(tab, "title", details.title || null);
+          pageAction.setProperty(tab, "title", details.title || null);
         },
 
         getTitle(details) {
           let tab = tabTracker.getTab(details.tabId);
 
-          let title = PageAction.for(extension).getProperty(tab, "title");
+          let title = pageAction.getProperty(tab, "title");
           return Promise.resolve(title);
         },
 
         setIcon(details) {
           let tab = tabTracker.getTab(details.tabId);
 
           let icon = IconDetails.normalize(details, extension, context);
-          PageAction.for(extension).setProperty(tab, "icon", icon);
+          pageAction.setProperty(tab, "icon", icon);
         },
 
         setPopup(details) {
           let tab = tabTracker.getTab(details.tabId);
 
           // Note: Chrome resolves arguments to setIcon relative to the calling
           // context, but resolves arguments to setPopup relative to the extension
           // root.
           // For internal consistency, we currently resolve both relative to the
           // calling context.
           let url = details.popup && context.uri.resolve(details.popup);
-          PageAction.for(extension).setProperty(tab, "popup", url);
+          pageAction.setProperty(tab, "popup", url);
         },
 
         getPopup(details) {
           let tab = tabTracker.getTab(details.tabId);
 
-          let popup = PageAction.for(extension).getProperty(tab, "popup");
+          let popup = pageAction.getProperty(tab, "popup");
           return Promise.resolve(popup);
         },
       },
     };
   }
 };
--- a/browser/components/extensions/ext-sidebarAction.js
+++ b/browser/components/extensions/ext-sidebarAction.js
@@ -286,35 +286,32 @@ extensions.on("ready", (type, extension)
 });
 /* eslint-enable mozilla/balanced-listeners */
 
 this.sidebarAction = class extends ExtensionAPI {
   onManifestEntry(entryName) {
     let {extension} = this;
     let {manifest} = extension;
 
-    let sidebarAction = new SidebarAction(manifest.sidebar_action, extension);
-    sidebarActionMap.set(extension, sidebarAction);
+    this.sidebarAction = new SidebarAction(manifest.sidebar_action, extension);
+    sidebarActionMap.set(extension, this.sidebarAction);
   }
 
   onShutdown(reason) {
-    let {extension} = this;
-
-    if (sidebarActionMap.has(extension)) {
-      // Don't remove everything on app shutdown so session restore can handle
-      // restoring open sidebars.
-      if (extension.shutdownReason !== "APP_SHUTDOWN") {
-        sidebarActionMap.get(extension).shutdown();
-      }
-      sidebarActionMap.delete(extension);
+    // Don't remove everything on app shutdown so session restore can handle
+    // restoring open sidebars.
+    if (reason !== "APP_SHUTDOWN") {
+      this.sidebarAction.shutdown();
     }
+    sidebarActionMap.delete(this.extension);
   }
 
   getAPI(context) {
     let {extension} = context;
+    const {sidebarAction} = this;
 
     function getTab(tabId) {
       if (tabId !== null) {
         return tabTracker.getTab(tabId);
       }
       return null;
     }
 
@@ -323,51 +320,51 @@ this.sidebarAction = class extends Exten
         async setTitle(details) {
           let nativeTab = getTab(details.tabId);
 
           let title = details.title;
           // Clear the tab-specific title when given a null string.
           if (nativeTab && title === "") {
             title = null;
           }
-          SidebarAction.for(extension).setProperty(nativeTab, "title", title);
+          sidebarAction.setProperty(nativeTab, "title", title);
         },
 
         getTitle(details) {
           let nativeTab = getTab(details.tabId);
 
-          let title = SidebarAction.for(extension).getProperty(nativeTab, "title");
+          let title = sidebarAction.getProperty(nativeTab, "title");
           return Promise.resolve(title);
         },
 
         async setIcon(details) {
           let nativeTab = getTab(details.tabId);
 
           let icon = IconDetails.normalize(details, extension, context);
-          SidebarAction.for(extension).setProperty(nativeTab, "icon", icon);
+          sidebarAction.setProperty(nativeTab, "icon", icon);
         },
 
         async setPanel(details) {
           let nativeTab = getTab(details.tabId);
 
           let url;
           // Clear the tab-specific url when given a null string.
           if (nativeTab && details.panel === "") {
             url = null;
           } else if (details.panel !== "") {
             url = context.uri.resolve(details.panel);
           } else {
             throw new ExtensionError("Invalid url for sidebar panel.");
           }
 
-          SidebarAction.for(extension).setProperty(nativeTab, "panel", url);
+          sidebarAction.setProperty(nativeTab, "panel", url);
         },
 
         getPanel(details) {
           let nativeTab = getTab(details.tabId);
 
-          let panel = SidebarAction.for(extension).getProperty(nativeTab, "panel");
+          let panel = sidebarAction.getProperty(nativeTab, "panel");
           return Promise.resolve(panel);
         },
       },
     };
   }
 };
--- a/toolkit/components/extensions/ext-backgroundPage.js
+++ b/toolkit/components/extensions/ext-backgroundPage.js
@@ -7,19 +7,16 @@ XPCOMUtils.defineLazyModuleGetter(this, 
                                   "resource://gre/modules/AddonManager.jsm");
 
 Cu.import("resource://gre/modules/ExtensionParent.jsm");
 var {
   HiddenExtensionPage,
   promiseExtensionViewLoaded,
 } = ExtensionParent;
 
-// WeakMap[Extension -> BackgroundPage]
-let backgroundPagesMap = new WeakMap();
-
 // Responsible for the background_page section of the manifest.
 class BackgroundPage extends HiddenExtensionPage {
   constructor(extension, options) {
     super(extension, "background");
 
     this.page = options.page || null;
     this.isGenerated = !!options.scripts;
     this.webNav = null;
@@ -74,26 +71,19 @@ class BackgroundPage extends HiddenExten
     }
 
     super.shutdown();
   }
 }
 
 this.backgroundPage = class extends ExtensionAPI {
   onManifestEntry(entryName) {
-    let {extension} = this;
-    let {manifest} = extension;
+    let {manifest} = this.extension;
 
-    let bgPage = new BackgroundPage(extension, manifest.background);
+    this.bgPage = new BackgroundPage(this.extension, manifest.background);
 
-    backgroundPagesMap.set(extension, bgPage);
-    return bgPage.build();
+    return this.bgPage.build();
   }
 
   onShutdown() {
-    let {extension} = this;
-
-    if (backgroundPagesMap.has(extension)) {
-      backgroundPagesMap.get(extension).shutdown();
-      backgroundPagesMap.delete(extension);
-    }
+    this.bgPage.shutdown();
   }
 };
--- a/toolkit/components/extensions/ext-geolocation.js
+++ b/toolkit/components/extensions/ext-geolocation.js
@@ -17,14 +17,15 @@ this.geolocation = class extends Extensi
       Services.perms.add(extension.principal.URI, "geo",
                          Services.perms.ALLOW_ACTION,
                          Services.perms.EXPIRE_SESSION);
     }
   }
 
   onShutdown() {
     let {extension} = this;
+
     if (extension.hasPermission("geolocation") &&
         Services.perms.testPermission(extension.principal.URI, "geo") == Services.perms.ALLOW_ACTION) {
       Services.perms.remove(extension.principal.URI, "geo");
     }
   }
 };
--- a/toolkit/components/extensions/ext-protocolHandlers.js
+++ b/toolkit/components/extensions/ext-protocolHandlers.js
@@ -5,18 +5,16 @@
 XPCOMUtils.defineLazyServiceGetter(this, "handlerService",
                                    "@mozilla.org/uriloader/handler-service;1",
                                    "nsIHandlerService");
 XPCOMUtils.defineLazyServiceGetter(this, "protocolService",
                                    "@mozilla.org/uriloader/external-protocol-service;1",
                                    "nsIExternalProtocolService");
 Cu.importGlobalProperties(["URL"]);
 
-const handlers = new WeakMap();
-
 function hasHandlerApp(handlerConfig) {
   let protoInfo = protocolService.getProtocolHandlerInfo(handlerConfig.protocol);
   let appHandlers = protoInfo.possibleApplicationHandlers;
   for (let i = 0; i < appHandlers.length; i++) {
     let handler = appHandlers.queryElementAt(i, Ci.nsISupports);
     if (handler instanceof Ci.nsIWebHandlerApp &&
         handler.uriTemplate === handlerConfig.uriTemplate) {
       return true;
@@ -39,37 +37,37 @@ this.protocolHandlers = class extends Ex
                       .createInstance(Ci.nsIWebHandlerApp);
       handler.name = handlerConfig.name;
       handler.uriTemplate = handlerConfig.uriTemplate;
 
       let protoInfo = protocolService.getProtocolHandlerInfo(handlerConfig.protocol);
       protoInfo.possibleApplicationHandlers.appendElement(handler, false);
       handlerService.store(protoInfo);
     }
-    handlers.set(extension, manifest.protocol_handlers);
   }
 
-  onShutdown() {
+  onShutdown(shutdownReason) {
     let {extension} = this;
+    let {manifest} = extension;
 
-    if (!handlers.has(extension) || extension.shutdownReason === "APP_SHUTDOWN") {
+    if (shutdownReason === "APP_SHUTDOWN") {
       return;
     }
-    for (let handlerConfig of handlers.get(extension)) {
+
+    for (let handlerConfig of manifest.protocol_handlers) {
       let protoInfo = protocolService.getProtocolHandlerInfo(handlerConfig.protocol);
       let appHandlers = protoInfo.possibleApplicationHandlers;
       for (let i = 0; i < appHandlers.length; i++) {
         let handler = appHandlers.queryElementAt(i, Ci.nsISupports);
         if (handler instanceof Ci.nsIWebHandlerApp &&
             handler.uriTemplate === handlerConfig.uriTemplate) {
           appHandlers.removeElementAt(i);
           if (protoInfo.preferredApplicationHandler === handler) {
             protoInfo.preferredApplicationHandler = null;
             protoInfo.alwaysAskBeforeHandling = true;
           }
           handlerService.store(protoInfo);
           break;
         }
       }
     }
-    handlers.delete(extension);
   }
 };
--- a/toolkit/components/extensions/ext-theme.js
+++ b/toolkit/components/extensions/ext-theme.js
@@ -6,19 +6,16 @@ XPCOMUtils.defineLazyModuleGetter(this, 
                                   "resource://gre/modules/Preferences.jsm");
 XPCOMUtils.defineLazyModuleGetter(this, "LightweightThemeManager",
                                   "resource://gre/modules/LightweightThemeManager.jsm");
 
 XPCOMUtils.defineLazyGetter(this, "gThemesEnabled", () => {
   return Preferences.get("extensions.webextensions.themes.enabled");
 });
 
-// WeakMap[Extension -> Theme]
-let themeMap = new WeakMap();
-
 const ICONS = Preferences.get("extensions.webextensions.themes.icons.buttons", "").split(",");
 
 /** Class representing a theme. */
 class Theme {
   /**
    * Creates a theme instance.
    *
    * @param {string} baseURI The base URI of the extension, used to
@@ -166,52 +163,42 @@ this.theme = class extends ExtensionAPI 
     let {extension} = this;
     let {manifest} = extension;
 
     if (!gThemesEnabled) {
       // Return early if themes are disabled.
       return;
     }
 
-    let theme = new Theme(extension.baseURI);
-    theme.load(manifest.theme);
-    themeMap.set(extension, theme);
+    this.theme = new Theme(extension.baseURI);
+    this.theme.load(manifest.theme);
   }
 
   onShutdown() {
-    let {extension} = this;
-
-    let theme = themeMap.get(extension);
-
-    if (!theme) {
-      // We won't have a theme if themes are disabled.
-      return;
+    if (this.theme) {
+      this.theme.unload();
     }
-
-    theme.unload();
   }
 
   getAPI(context) {
     let {extension} = context;
+
     return {
       theme: {
-        update(details) {
+        update: (details) => {
           if (!gThemesEnabled) {
             // Return early if themes are disabled.
             return;
           }
 
-          let theme = themeMap.get(extension);
-
-          if (!theme) {
+          if (!this.theme) {
             // WebExtensions using the Theme API will not have a theme defined
             // in the manifest. Therefore, we need to initialize the theme the
             // first time browser.theme.update is called.
-            theme = new Theme(extension.baseURI);
-            themeMap.set(extension, theme);
+            this.theme = new Theme(extension.baseURI);
           }
 
-          theme.load(details);
+          this.theme.load(details);
         },
       },
     };
   }
 };