WorkerDOM.upgrade should wait for Worker init. (#1047)

samouri · web-flow · commit 1472534283a7 · 2021-04-20T14:44:07.000-04:00
* Ensure WorkerDOM.upgrade() waits for iframe init complete

* format

* expand sandboxed example to cover callFunction.

* fix test

* extra wiggleroom
diff --git a/demo/sandboxed/index.html b/demo/sandboxed/index.html
@@ -7,26 +7,28 @@
     <script src="/dist/amp-debug/main.mjs" type="module"></script>
   </head>
   <body>
-    <div src="/hello-world/hello-world.js" id="upgrade-me">
+    <div src="/sandboxed/index.js" id="upgrade-me">
       <div class="root">
         <button>Insert Hello World!</button>
       </div>
     </div>
     <script type="module">
       import { upgradeElement } from '/dist/amp-debug/main.mjs';
-      upgradeElement(
+      const worker = await upgradeElement(
         document.getElementById('upgrade-me'),
         '/dist/amp-debug/worker/worker.mjs',
         () => {},
         undefined,
         { iframeUrl: 'sandbox-iframe.html' }
-      ).then((worker) => {
-        worker.onmessage = (msg) =>
-          console.log(`onmessage: ${JSON.stringify(msg)}`);
-        worker.onmessageerror = (msg) =>
-          console.error(`msgerror ${JSON.stringify(msg)}`);
-        worker.onerror = (msg) =>
-          console.error(`error: ${JSON.stringify(msg)}`);
+      );
+      worker.onmessage = (msg) =>
+        console.log(`onmessage: ${JSON.stringify(msg)}`);
+      worker.onmessageerror = (msg) =>
+        console.error(`msgerror ${JSON.stringify(msg)}`);
+      worker.onerror = (msg) => console.error(`error: ${JSON.stringify(msg)}`);
+
+      worker.callFunction('getData').then((data) => {
+        console.log(`Received data: ${JSON.stringify(data)}`);
       });
     </script>
   </body>
diff --git a/demo/sandboxed/sandbox-iframe.html b/demo/sandboxed/sandbox-iframe.html
@@ -7,6 +7,7 @@
     const MESSAGE_TYPES = {
       ready: 'iframe-ready',
       init: 'init-worker',
+      workerReady: 'worker-ready',
       onmessage: 'onmessage',
       onerror: 'onerror',
       onmessageerror: 'onmessageerror',
@@ -58,6 +59,8 @@
       listen(MESSAGE_TYPES.postMessage, ({ message }) =>
         worker.postMessage(message)
       );
+
+      send(MESSAGE_TYPES.workerReady);
     });
   </script>
 </html>
diff --git a/package.json b/package.json
@@ -99,7 +99,7 @@
       "brotli": "12 kB"
     },
     "./dist/worker/worker.js": {
-      "brotli": "13.5 kB"
+      "brotli": "14 kB"
     },
     "./dist/amp-production/main.mjs": {
       "brotli": "4.5 kB"
diff --git a/src/main-thread/iframe-worker.ts b/src/main-thread/iframe-worker.ts
@@ -18,7 +18,7 @@ type MessageFromWorker = {
   type: 'onmessage' | 'onerror' | 'onmessageerror';
   message: any;
 };
-export type MessageFromIframe = { type: 'iframe-ready' } | MessageFromWorker;
+export type MessageFromIframe = { type: 'iframe-ready' } | { type: 'worker-ready' } | MessageFromWorker;
 export type MessageToIframe = { type: 'terminate' } | { type: 'init-worker'; code: string } | { type: 'postMessage'; message: any };
 
 /**
@@ -29,16 +29,19 @@ export type MessageToIframe = { type: 'terminate' } | { type: 'init-worker'; cod
  * The iframe used for sandboxing must follow a specific contract. It:
  *   1. Must send a ready message to the main-thread.
  *   2. Must listen for a message from main-thread with the code to initialize a Worker with.
- *   3. Must proxy all messages between the Worker and Parent, including errors.
+ *   3. Must send "worker-ready" once worker is initialized.
+ *   4. Must proxy all messages between the Worker and Parent, including errors.
  */
 class IframeWorker {
   // Public Worker API
   public onerror: (this: IframeWorker, ev: ErrorEvent) => any;
   public onmessage: (this: IframeWorker, ev: MessageEvent) => any;
   public onmessageerror: (this: IframeWorker, ev: MessageEvent) => any;
+  public readyPromise: Promise<void>;
 
   // Internal variables.
   private iframe: HTMLIFrameElement;
+  private readyPromiseResolve: Function;
 
   /**
    * @param url The URL to initiate the worker from.
@@ -50,6 +53,9 @@ class IframeWorker {
     this.iframe.setAttribute('style', 'display:none');
     this.iframe.setAttribute('src', iframeUrl);
     this.url = url;
+    this.readyPromise = new Promise((resolve) => {
+      this.readyPromiseResolve = resolve;
+    });
 
     this.setupInit();
     this.proxyFromWorker();
@@ -65,11 +71,14 @@ class IframeWorker {
       fetch(this.url.toString())
         .then((res) => res.text())
         .then((code) => {
-          if ((event.data as MessageFromIframe).type == 'iframe-ready') {
+          const data = event.data as MessageFromIframe;
+          if (data.type == 'iframe-ready') {
             const msg: MessageToIframe = { type: 'init-worker', code };
             this.iframe.contentWindow!.postMessage(msg, '*');
+          } else if (data.type === 'worker-ready') {
+            this.readyPromiseResolve();
+            window.removeEventListener('message', listener);
           }
-          window.removeEventListener('message', listener);
         });
     };
     window.addEventListener('message', listener);
@@ -99,7 +108,9 @@ class IframeWorker {
    */
   postMessage(message: any, transferables?: Array<Transferable>) {
     const msg: MessageToIframe = { type: 'postMessage', message };
-    this.iframe.contentWindow!.postMessage(msg, '*', transferables);
+    this.readyPromise.then(() => {
+      this.iframe.contentWindow!.postMessage(msg, '*', transferables);
+    });
   }
 
   /**
diff --git a/src/main-thread/install.ts b/src/main-thread/install.ts
@@ -80,7 +80,7 @@ export function install(
         }
       };
 
-      return new ExportedWorker(workerContext, normalizedConfig);
+      return workerContext.ready().then(() => new ExportedWorker(workerContext, normalizedConfig));
     }
     return null;
   });
diff --git a/src/main-thread/worker.ts b/src/main-thread/worker.ts
@@ -94,6 +94,14 @@ export class WorkerContext {
     }
   }
 
+  /**
+   * Returns a Promise that resolves when the Worker is ready to receive messages.
+   * @returns {Promise<void>}
+   */
+  ready() {
+    return (this.worker as IframeWorker).readyPromise || Promise.resolve();
+  }
+
   /**
    * Returns the private worker.
    */
diff --git a/src/test/main-thread/iframe-worker.test.ts b/src/test/main-thread/iframe-worker.test.ts
@@ -78,10 +78,14 @@ test('Should listen for and respond to an iframe-ready signal.', async (t) => {
   ]);
 });
 
-test('Should proxy postMessage calls.', async (t) => {
-  const { worker, sentToIframe } = t.context;
+test('Should proxy postMessage calls after init completes.', async (t) => {
+  const { worker, sentToIframe, fakeReceiveMessage } = t.context;
+
   worker.postMessage({ hello: 'world' });
+  t.is(sentToIframe.length, 0);
 
+  fakeReceiveMessage({ type: 'worker-ready' });
+  await new Promise(setTimeout as any);
   t.is(sentToIframe.length, 1);
   t.deepEqual(sentToIframe, [
     {
@@ -101,14 +105,11 @@ test('Should proxy iframed worker messages.', async (t) => {
   worker.onmessageerror = (msg: MessageEvent) => onmessageerrorMessages.push(msg);
   worker.onerror = (msg: ErrorEvent) => onerrorMessages.push(msg);
 
-  fakeReceiveMessage({ type: 'onmessage', message: { onmessage: 1 } });
-  fakeReceiveMessage({ type: 'onmessage', message: { onmessage: 2 } });
+  fakeReceiveMessage({ type: 'onmessage', message: {} });
+  fakeReceiveMessage({ type: 'onmessage', message: {} });
 
   fakeReceiveMessage({ type: 'onerror', message: { onerror: 1 } });
-  fakeReceiveMessage({
-    type: 'onmessageerror',
-    message: { onmessageerror: 1 },
-  });
+  fakeReceiveMessage({ type: 'onmessageerror', message: {} });
 
   t.is(onmessageMessages.length, 2);
   t.is(onmessageerrorMessages.length, 1);

Original file line number	Diff line number	Diff line change
`@@ -80,7 +80,7 @@ export function install(`
`80`	`80`	`}`
`81`	`81`	`};`
`82`	`82`
`83`		`- return new ExportedWorker(workerContext, normalizedConfig);`
	`83`	`+ return workerContext.ready().then(() => new ExportedWorker(workerContext, normalizedConfig));`
`84`	`84`	`}`
`85`	`85`	`return null;`
`86`	`86`	`});`
Original file line number	Diff line number	Diff line change
`@@ -94,6 +94,14 @@ export class WorkerContext {`
`94`	`94`	`}`
`95`	`95`	`}`
`96`	`96`
	`97`	`+ /**`
	`98`	`+ * Returns a Promise that resolves when the Worker is ready to receive messages.`
	`99`	`+ * @returns {Promise<void>}`
	`100`	`+ */`
	`101`	`+ ready() {`
	`102`	`+ return (this.worker as IframeWorker).readyPromise \|\| Promise.resolve();`
	`103`	`+ }`
	`104`	`+`
`97`	`105`	`/**`
`98`	`106`	`* Returns the private worker.`
`99`	`107`	`*/`