mikehall314 / electron-comlink

An adapter for Electron IPC to allow communication with Comlink
MIT License
8 stars 1 forks source link

TransferList Experimentation #1

Open kgullion opened 3 years ago

kgullion commented 3 years ago

Comlink uses MessagePorts internally for proxy endpoints and Electron has MessageChannel support via MessageChannelMain now.

By patching the MessagePortMain, it can now be used in main.js as a Comlink endpoint. This seems to work pretty well in my limited testing (example below).

Unfortunately, I've been unable to get proxy working. I was hoping replacing the proxyTransferHandler would be enough but while the MessagePorts are transferable, it seems they are not StructuredCloneable :(

I'm guessing it would be possible to get proxy working if we had access to the transferred MessagePort itself in deserialize. Not sure where to go from here but it seems potentially doable. Thoughts?

// main.js

const proxyTransferHandler = {
  canHandle: (val) => !!val && val[comlink.proxyMarker],
  serialize(obj) {
    const { port1, port2 } = new MessageChannelMain();
    console.log(obj, port1, port2);
    comlink.expose(obj, patchMessagePort(port1));
    return [port2, [port2]];
  },
  deserialize(port) {
    console.log(port);
    port.start();
    return comlink.wrap(patchMessagePort(port));
  },
};

comlink.transferHandlers.set("proxy", proxyTransferHandler);

function patchMessagePort(port) {
  if (!port.addEventListener) port.addEventListener = port.on;
  if (!port.removeEventListener) port.removeEventListener = port.off;
  return port;
}

ipcMain.on("comlink1", ({ ports }) => {
  const port = ports[0];
  const obj = { log: console.log, inc: (a) => a + 1, callWith1: (cb) => cb(1) };
  comlink.expose(obj, patchMessagePort(port));
});

ipcMain.on("comlink2", ({ ports }) => {
  const port = ports[0];
  const remote = comlink.wrap(patchMessagePort(port));
  remote.log("hello from main");
  console.log("4+1=", remote.inc(4));
  // console.log('cb', remote.callWith2(comlink.proxy((a)=>a+1)))
});
// preload.js

function comlinkExample1() {
  const { port1, port2 } = new MessageChannel();
  ipcRenderer.postMessage("comlink1", null, [port2]);

  const remote = wrap(port1);
  remote.log("hello from preload");
  console.log("2+1=", remote.inc(2));
  // console.log('cb', remote.callWith1(proxy((a)=>a+1)))
}
comlinkExample1();

function comlinkExample2() {
  const { port1, port2 } = new MessageChannel();
  ipcRenderer.postMessage("comlink2", null, [port2]);

  const obj = { log: console.log, inc: (a) => a + 1, callWith2: (cb) => cb(2) };
  expose(obj, port1);
}
comlinkExample2();
kgullion commented 3 years ago

Came up with a proof of concept for passing proxy ports via a backchannel.

Two big issues with this trick:

  1. I believe this will only work for the first ipcRenderer though due to how the new proxyTransferHandler works.
  2. The proxied function comes out the other end as a Promise instead of a callable. Easy enough to just await but it means this doesn't behave the same as the comlink API.

Issue 1 seems fixable but I don't see an easy way around issue 2 unless there is a way to do requestEndpoint synchronously. This meets my current usecase (whole thing came out of trying to implement this) so I probably won't spent too much more time on it but I hope it's at least somewhat useful.

const comlink = require("comlink");

// from https://github.com/GoogleChromeLabs/comlink/blob/v4.3.0/src/comlink.ts#L546
function generateUUID() {
  return new Array(4)
    .fill(0)
    .map(() => Math.floor(Math.random() * Number.MAX_SAFE_INTEGER).toString(16))
    .join("-");
}

// from https://github.com/GoogleChromeLabs/comlink/blob/v4.3.0/src/comlink.ts#L180
const isObject = (val) =>
  (typeof val === "object" && val !== null) || typeof val === "function";

// patch electron MessagePortMain to conform to comlink Endpoint schema
function patchMessagePort(port) {
  if (!port.addEventListener) port.addEventListener = port.on;
  if (!port.removeEventListener) port.removeEventListener = port.off;
  return port;
}

// object store holds object until an endpoint request is received
function createEndpointStore(proxyPort, MessageChannelClass) {
  const store = new Map();
  const onEndpointRequest = (e) => {
    // retrieve serialized object
    const id = e.data;
    const obj = store.get(id);
    // create comlink endpoints
    const { port1, port2 } = new MessageChannelClass();
    // expose serialized object on port1
    comlink.expose(obj, port1);
    // send port2 back to caller
    e.ports[0].postMessage(id, [port2]);
    // delete obj from store?
    // store.delete(id) // is this needed?
  };

  // endpoint requests will arrive on the passed proxy port
  const prox = proxyPort;
  prox.addEventListener("message", onEndpointRequest);
  if (prox.start) prox.start();

  return {
    saveObject(obj) {
      // store the object and send an id to the caller in place of an endpoint
      const id = generateUUID();
      store.set(id, obj);
      return id;
    },
    requestEndpoint(id) {
      // create channel to request endpoint
      const { port1, port2 } = new MessageChannelClass();
      // send request
      prox.postMessage(id, [port2]);
      // return endpoint via promise
      return new Promise((resolve) => {
        port1.addEventListener("message", (e) => {
          if (e.data === id) resolve(patchMessagePort(e.ports[0]));
        });
        port1.start && port1.start();
      });
    },
  };
}

// MessagePortMain is not StructuredCloneable but the default comlink proxy transferHandler
// relies on being able to pass a MessagePort via the postMessage body.
// Instead, we patch the transferHandler to pass the ports via a backchannel
// patches https://github.com/GoogleChromeLabs/comlink/blob/v4.3.0/src/comlink.ts#L215
function patchProxyTransferHandler(proxyPort, MessageChannelClass) {
  // endpoint store is used to transfer ports via backchannel
  const { saveObject, requestEndpoint } = createEndpointStore(
    proxyPort,
    MessageChannelClass
  );
  // remotes stores the wrapped endpoint after first request
  const remotes = new Map();
  // new transferHandler deserializes to a Promise rather than a function
  const proxyTransferHandler = {
    canHandle: (val) => isObject(val) && val[comlink.proxyMarker],
    serialize: (obj) => {
      return [saveObject(obj), []];
    },
    deserialize: async (id) => {
      if (remotes.has(id)) return remotes.get(id);
      // get endpoint, setup remote, and save for later
      const ep = await requestEndpoint(id);
      const remote = comlink.wrap(ep);
      remotes.set(id, remote);
      if (ep.start) ep.start();
      return remote;
    },
  };
  // proxy is finally patched here
  comlink.transferHandlers.set("proxy", proxyTransferHandler);
}

// make sure we only patch the proxy one time for each environment
let proxyIsPatched = false;

function patchMainProxy(ipc) {
  if (proxyIsPatched) return;
  proxyIsPatched = true;
  // electron doesn't have MessageChannel, use a patched MessageChannelMain instead
  const { MessageChannelMain } = require("electron");
  class MessageChannelPatched {
    constructor() {
      const { port1, port2 } = new MessageChannelMain();
      this.port1 = patchMessagePort(port1);
      this.port2 = patchMessagePort(port2);
    }
  }
  ipc.on("comlink-prox-port", (e) =>
    patchProxyTransferHandler(
      patchMessagePort(e.ports[0]),
      MessageChannelPatched
    )
  );
}

function patchRendererProxy(ipc) {
  if (proxyIsPatched) return;
  proxyIsPatched = true;
  const { port1, port2 } = new MessageChannel();
  ipc.postMessage("comlink-prox-port", null, [port2]);
  patchProxyTransferHandler(port1, MessageChannel);
}

// waits to receive a MessagePort from renderer
function ipcMainEndpoint(ipc, name = "") {
  const id = "comlink-port-init-" + name;
  patchMainProxy(ipc);
  return new Promise((resolve) =>
    ipc.on(id, (e) => resolve(patchMessagePort(e.ports[0])))
  );
}

// creates both endpoint ports, then sends one to the backend
function ipcRendererEndpoint(ipc, name = "") {
  const id = "comlink-port-init-" + name;
  patchRendererProxy(ipc);
  const { port1, port2 } = new MessageChannel();
  ipc.postMessage(id, null, [port2]);
  return Promise.resolve(port1);
}

exports.ipcMainEndpoint = ipcMainEndpoint;
exports.ipcRendererEndpoint = ipcRendererEndpoint;
mikehall314 commented 3 years ago

This looks awesome -- I'm crazy busy starting a new job at the moment, but I'll take a look as soon as I can.

kgullion commented 3 years ago

Congrats on the new job! :tada:

I've uploaded a full proof of concept here: https://github.com/kgullion/comlink-electron-proof-of-concept

The repo was created from this template purely out of convenience, sorry for all the extra cruft.

The logic is spread across three files in main, preload, and renderer

There is a lot of duplication between main and renderer, along with subtle differences between the two (mostly on vs addEventListener). main also includes two patch functions for the backend proxyTransferHandler.

Proxy behaves correctly now, resolving to a callable instead of a promise. To use it, one of the patch functions must be called to fix the backend logic. The frontend proxy handler doesn't need patched meaning it should still work as expected for WebWorkers and the like.

The way I'm setting up the initial port transfer seems a bit overkill, probably don't need to even use MessagePorts now that I'm thinking about it. It's also still a bit wonky if there is a channelname collision between two browsers. Bit out of scope for this library but worth a mention.

There is example code in main and renderer, the runAll function for each is called in the respective index.ts.