plugin: message channels for dynamic plugins

Author: psybuzz

tensorboard/components/plugin_util/message.ts

@@ -13,6 +13,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+/**
+ * This file defines utilities shared by TensorBoard (plugin host) and the
+ * dynamic plugin library, used by plugin authors.
+ */
+
 export type PayloadType =
   | null
   | undefined
@@ -30,6 +35,7 @@ export interface Message {
   id: string;
   payload: PayloadType;
   error: string | null;
+  isReply: boolean
 }
 
 export type MessageType = string;
@@ -40,21 +46,15 @@ interface PromiseResolver {
   reject: (error: Error) => void;
 }
 
-export abstract class IPC {
-  private idPrefix: string;
+export class IPC {
   private id = 0;
   private readonly responseWaits = new Map<string, PromiseResolver>();
   private readonly listeners = new Map<MessageType, MessageCallback>();
+  private readonly port: MessagePort;
 
-  constructor() {
-    window.addEventListener('message', this.onMessage.bind(this));
-
-    // TODO(tensorboard-team): remove this by using MessageChannel.
-    const randomArray = new Uint8Array(16);
-    window.crypto.getRandomValues(randomArray);
-    this.idPrefix = Array.from(randomArray)
-      .map((int: number) => int.toString(16))
-      .join('');
+  constructor(port) {
+    this.port = port;
+    port.addEventListener('message', this.onMessage.bind(this));
   }
 
   listen(type: MessageType, callback: MessageCallback) {
@@ -66,13 +66,12 @@ export abstract class IPC {
   }
 
   private async onMessage(event: MessageEvent) {
-    // There are instances where random browser extensions send messages.
-    if (typeof event.data !== 'string') return;
-
     const message = JSON.parse(event.data) as Message;
     const callback = this.listeners.get(message.type);
 
-    if (this.responseWaits.has(message.id)) {
+    if (message.isReply) {
+      if (!this.responseWaits.has(message.id))
+        return;
       const {id, payload, error} = message;
       const {resolve, reject} = this.responseWaits.get(id);
       this.responseWaits.delete(id);
@@ -100,22 +99,22 @@ export abstract class IPC {
       id: message.id,
       payload,
       error,
+      isReply: true,
     };
-    this.postMessage(event.source, JSON.stringify(replyMessage));
+    this.postMessage(replyMessage);
   }
 
-  private postMessage(targetWindow: Window, message: string) {
-    targetWindow.postMessage(message, '*');
+  private postMessage(message: Message) {
+    this.port.postMessage(JSON.stringify(message));
   }
 
-  protected sendMessageToWindow(
-    targetWindow: Window,
+  sendMessage(
     type: MessageType,
-    payload: PayloadType
+    payload: PayloadType,
   ): Promise<PayloadType> {
-    const id = `${this.idPrefix}_${this.id++}`;
-    const message: Message = {type, id, payload, error: null};
-    this.postMessage(targetWindow, JSON.stringify(message));
+    const id = `${this.id++}`;
+    const message: Message = {type, id, payload, error: null, isReply: false};
+    this.postMessage(message);
     return new Promise((resolve, reject) => {
       this.responseWaits.set(id, {resolve, reject});
     });

tensorboard/components/plugin_util/plugin-guest.ts

@@ -12,19 +12,25 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-import {IPC, Message, MessageType, PayloadType} from './message.js';
-
-class GuestIPC extends IPC {
-  /**
-   * payload must be JSON serializable.
-   */
-  sendMessage(type: MessageType, payload: PayloadType): Promise<PayloadType> {
-    return this.sendMessageToWindow(window.parent, type, payload);
-  }
+import {IPC, Message} from './message.js';
+
+/**
+ * This code is part of a public bundle provided to plugin authors,
+ * and runs within an IFrame to setup communication with TensorBoard's frame.
+ */
+if (!window.parent) {
+  throw Error('This library must be run from within a loaded TensorBoard dynamic plugin.');
 }
 
+const channel = new MessageChannel();
+const ipc = new IPC(channel.port1);
+channel.port1.start();
+
+const VERSION = 'experimental';
+window.parent.postMessage(`${VERSION}.bootstrap`, '*', [channel.port2]);
+
 // Only export for testability.
-export const _guestIPC = new GuestIPC();
+export const _guestIPC = ipc;
 
 /**
  * Sends a message to the parent frame.
@@ -45,3 +51,7 @@ export const listen = _guestIPC.listen.bind(_guestIPC);
  * Unsubscribes a callback to a message.
  */
 export const unlisten = _guestIPC.unlisten.bind(_guestIPC);
+
+
+// TODO: Register API for authors.
+// Methods should use 'ipc.sendMessage', Events should use 'ipc.listen'.

tensorboard/components/plugin_util/plugin-host.ts

@@ -12,46 +12,63 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-import {IPC, Message, MessageType, PayloadType} from './message.js';
-
-class HostIPC extends IPC {
-  sendMessage(
-    iframe: HTMLIFrameElement,
-    type: MessageType,
-    payload: PayloadType
-  ): Promise<PayloadType> {
-    return this.sendMessageToWindow(iframe.contentWindow, type, payload);
-  }
-}
+import {IPC, MessageType, PayloadType} from './message.js';
+
+const portIPCs = new Set<IPC>();
+const VERSION = 'experimental';
+const ipcToFrame = new WeakMap<IPC, HTMLIFrameElement>();
+
+// The initial Window-level listener is needed to bootstrap only.
+// All further communication is done over MessagePorts.
+window.addEventListener('message', (event) => {
+  if (event.data !== `${VERSION}.bootstrap`)
+    return;
+  const port = event.ports[0];
+  if (!port)
+    return;
+  const frame = event.source ? event.source.frameElement : null;
+  if (!frame)
+    return;
 
-const hostIPC = new HostIPC();
-const _listen = hostIPC.listen.bind(hostIPC);
-const _unlisten = hostIPC.unlisten.bind(hostIPC);
-const _sendMessage = hostIPC.sendMessage.bind(hostIPC);
+  const portIPC = new IPC(port);
+  portIPCs.add(portIPC);
+  ipcToFrame.set(portIPC, frame as HTMLIFrameElement);
+  port.start();
+
+  // TODO: install API.
+});
+
+function _broadcast(
+  type: MessageType,
+  payload: PayloadType
+): Promise<PayloadType[]> {
+  // Clean up disconnected iframes, since they won't respond.
+  for (const ipc of portIPCs) {
+    if (!ipcToFrame.get(ipc).isConnected) {
+      portIPCs.delete(ipc);
+      ipcToFrame.delete(ipc);
+    }
+  }
 
-export const sendMessage = _sendMessage;
-export const listen = _listen;
-export const unlisten = _unlisten;
+  const ipcs = [...portIPCs];
+  const promises = ipcs.map(ipc => ipc.sendMessage(type, payload));
+  return Promise.all(promises);
+}
 
-// Export for testability.
-export const _hostIPC = hostIPC;
+export const broadcast = _broadcast;
 
 namespace tf_plugin {
+
   /**
-   * Sends a message to the frame specified.
-   * @return Promise that resolves with a payload from frame in response to the message.
+   * Sends a message to all dynamic plugins. Individual plugins decide whether
+   * or not to listen.
+   * @return Promise that resolves with a list of payloads from each plugin's
+   *         response (or null) to the message.
    *
    * @example
-   * const someList = await sendMessage('v1.some.type.guest.understands');
+   * const someList = await broadcast('some.type.guest.understands');
    * // do fun things with someList.
    */
-  export const sendMessage = _sendMessage;
-  /**
-   * Subscribes to messages from specified frame of a type specified.
-   */
-  export const listen = _listen;
-  /**
-   * Unsubscribes to messages from specified frame of a type specified.
-   */
-  export const unlisten = _unlisten;
+  export const broadcast = _broadcast;
+
 } // namespace tf_plugin

tensorboard/components/plugin_util/test/plugin-test.ts

@@ -57,24 +57,15 @@ namespace tf_plugin.test {
           this.destListen = this.guestWindow.test.listen;
           this.destUnlisten = this.guestWindow.test.unlisten;
           this.srcSendMessage = (type, payload) => {
-            return pluginHost.sendMessage(this.guestFrame, type, payload);
+            return new Promise(async resolve => {
+              const results = await pluginHost.broadcast(type, payload);
+              resolve(results[0]);
+            });
           };
           this.destPostMessageSpy = () =>
             this.sandbox.spy(this.guestWindow.test._guestIPC, 'postMessage');
         },
       },
-      {
-        spec: 'guest (src) to host (dest)',
-        beforeEachFunc: function() {
-          this.destListen = pluginHost.listen;
-          this.destUnlisten = pluginHost.unlisten;
-          this.srcSendMessage = (type, payload) => {
-            return this.guestWindow.test.sendMessage(type, payload);
-          };
-          this.destPostMessageSpy = () =>
-            this.sandbox.spy(pluginHost._hostIPC, 'postMessage');
-        },
-      },
     ].forEach(({spec, beforeEachFunc}) => {
       describe(spec, () => {
         beforeEach(beforeEachFunc);