Bug 1728326 - Notify main process when the extension background service worker has been fully loaded. r=zombie

Additional changes needed to observe the "extension-worker-loaded" topics and notify the main process when the background
service worker script has been fully loaded (similarly to how we notify the main process when
the background page has been fully loaded).

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

toolkit/components/extensions/ExtensionChild.jsm

@@ -757,7 +757,13 @@ class ChildAPIManager {
 
     this.conduit = context.openConduit(this, {
       childId: this.id,
-      send: ["CreateProxyContext", "APICall", "AddListener", "RemoveListener"],
+      send: [
+        "CreateProxyContext",
+        "ContextLoaded",
+        "APICall",
+        "AddListener",
+        "RemoveListener",
+      ],
       recv: ["CallResult", "RunListener"],
     });
 

toolkit/components/extensions/ExtensionParent.jsm

@@ -746,7 +746,13 @@ ParentAPIManager = {
     this.conduit = new BroadcastConduit(this, {
       id: "ParentAPIManager",
       reportOnClosed: "childId",
-      recv: ["CreateProxyContext", "APICall", "AddListener", "RemoveListener"],
+      recv: [
+        "CreateProxyContext",
+        "ContextLoaded",
+        "APICall",
+        "AddListener",
+        "RemoveListener",
+      ],
       send: ["CallResult"],
       query: ["RunListener"],
     });
@@ -867,6 +873,13 @@ ParentAPIManager = {
     this.proxyContexts.set(childId, context);
   },
 
+  recvContextLoaded(data, { actor, sender }) {
+    let context = this.getContextById(data.childId);
+    verifyActorForContext(actor, context);
+    const { extension } = context;
+    extension.emit("extension-proxy-context-load:completed", context);
+  },
+
   recvConduitClosed(sender) {
     this.closeProxyContext(sender.id);
   },
@@ -1562,6 +1575,40 @@ function watchExtensionProxyContextLoad(
   };
 }
 
+/**
+ * This helper is used to subscribe a listener (e.g. in the ext-backgroundPage)
+ * to be called for every ExtensionProxyContext created for an extension
+ * background service worker given its related extension.
+ *
+ * @param {object} params.extension
+ *        The Extension on which we are going to listen for the newly created ExtensionProxyContext.
+ * @param {function} onExtensionWorkerContextLoaded
+ *        The callback that is called when the worker script has been fully loaded (as `callback(context)`);
+ *
+ * @returns {function}
+ *          Unsubscribe the listener.
+ */
+function watchExtensionWorkerContextLoaded(
+  { extension },
+  onExtensionWorkerContextLoaded
+) {
+  if (typeof onExtensionWorkerContextLoaded !== "function") {
+    throw new Error("Missing onExtensionWorkerContextLoaded handler");
+  }
+
+  const listener = (event, context) => {
+    if (context.viewType == "background_worker") {
+      onExtensionWorkerContextLoaded(context);
+    }
+  };
+
+  extension.on("extension-proxy-context-load:completed", listener);
+
+  return () => {
+    extension.off("extension-proxy-context-load:completed", listener);
+  };
+}
+
 // Manages icon details for toolbar buttons in the |pageAction| and
 // |browserAction| APIs.
 let IconDetails = {
@@ -1914,6 +1961,7 @@ var ExtensionParent = {
   apiManager,
   promiseExtensionViewLoaded,
   watchExtensionProxyContextLoad,
+  watchExtensionWorkerContextLoaded,
   DebugUtils,
 };
 

toolkit/components/extensions/ExtensionProcessScript.jsm

@@ -359,6 +359,13 @@ var ExtensionAPIRequestHandler = {
     );
   },
 
+  onExtensionWorkerLoaded(policy, serviceWorkerDescriptorId) {
+    ExtensionWorkerChild.notifyExtensionWorkerContextLoaded(
+      serviceWorkerDescriptorId,
+      policy
+    );
+  },
+
   onExtensionWorkerDestroyed(policy, serviceWorkerDescriptorId) {
     ExtensionWorkerChild.destroyExtensionWorkerContext(
       serviceWorkerDescriptorId

toolkit/components/extensions/ExtensionWorkerChild.jsm

@@ -635,6 +635,13 @@ class WorkerContextChild extends BaseContext {
     return conduit;
   }
 
+  notifyWorkerLoaded() {
+    this.childManager.conduit.sendContextLoaded({
+      childId: this.childManager.id,
+      extensionId: this.extension.id,
+    });
+  }
+
   withAPIRequest(request, callable) {
     this.webidlAPIRequest = request;
     try {
@@ -768,6 +775,27 @@ var ExtensionWorkerChild = {
     return null;
   },
 
+  /**
+   * Notify the main process when an extension worker script has been loaded.
+   *
+   * @param {number} descriptorId The service worker descriptor ID of the destroyed context.
+   * @param {WebExtensionPolicy} policy
+   */
+  notifyExtensionWorkerContextLoaded(descriptorId, policy) {
+    let context = this.extensionWorkerContexts.get(descriptorId);
+    if (context) {
+      if (context.extension.id !== policy.id) {
+        Cu.reportError(
+          new Error(
+            `ServiceWorker ${descriptorId} does not belong to the expected extension: ${policy.id}`
+          )
+        );
+        return;
+      }
+      context.notifyWorkerLoaded();
+    }
+  },
+
   /**
    * Close the ExtensionWorkerContextChild belonging to the given service worker, if any.
    *

toolkit/components/extensions/parent/ext-backgroundPage.js

@@ -7,7 +7,11 @@
 var { ExtensionParent } = ChromeUtils.import(
   "resource://gre/modules/ExtensionParent.jsm"
 );
-var { HiddenExtensionPage, promiseExtensionViewLoaded } = ExtensionParent;
+var {
+  HiddenExtensionPage,
+  promiseExtensionViewLoaded,
+  watchExtensionWorkerContextLoaded,
+} = ExtensionParent;
 
 ChromeUtils.defineModuleGetter(
   this,
@@ -123,7 +127,18 @@ class BackgroundWorker {
   async build() {
     const { extension } = this;
 
+    let context;
     try {
+      const contextPromise = new Promise(resolve => {
+        let unwatch = watchExtensionWorkerContextLoaded(
+          { extension, viewType: "background_worker" },
+          context => {
+            unwatch();
+            resolve(context);
+          }
+        );
+      });
+
       // TODO(Bug 17228327): follow up to spawn the active worker for a previously installed
       // background service worker.
       const regInfo = await serviceWorkerManager.registerForAddonPrincipal(
@@ -133,8 +148,8 @@ class BackgroundWorker {
         Ci.nsIServiceWorkerRegistrationInfo
       );
 
-      // TODO(bug 17228326): wait for worker context to be loaded (as we currently do
-      // for the delayed background page).
+      context = await contextPromise;
+
       await this.waitForActiveWorker();
     } catch (e) {
       // Extension may be shutting down before the background worker has registered or
@@ -150,10 +165,12 @@ class BackgroundWorker {
       return;
     }
 
-    // TODO(bug 17228326): wait for worker context to be loaded and
-    // wait for all persistent event listeners registered by the worker
-    // script to be handled (as we currently do for the delayed background
-    // page).
+    if (context) {
+      // Wait until all event listeners registered by the script so far
+      // to be handled.
+      await Promise.all(context.listenerPromises);
+      context.listenerPromises = null;
+    }
 
     if (extension.persistentListeners) {
       // |this.extension| may be null if the extension was shut down.