partytown の Worker からの同期的メインスレッド操作を実装してみる

Partytown とは

GitHub - BuilderIO/partytown: Relocate resource intensive third-party scripts off of the main thread and into a web worker. 🎉

今Partytownがヤバい。JavaScript Sandboxの未来はどっちだ？

要は 3rd party script を安全に隔離するため、 WebWorker + DOM の mock で動かす。

GitHub - ampproject/worker-dom: The same DOM API and Frameworks you know, but in a Web Worker.

この DOM がすごい2018: worker-dom - mizchi's blog

worker-dom との一番の違いは、worker-dom は DOM API をすべて mock 化して実装してその操作を転送する。その上で取得する操作は非同期APIとして実装される。

const style = await window.getComputedStyle(element)

partytown は非同期操作を隠蔽する

const style = window.getComputedStyle(element)

これがなぜ可能なのかを調べた。

Web Worker で main-thread 操作を同期的に行うテクニック

• main-thread から worker と service-worker を起動する
• worker から service-worker に同期 xhr で (blocking状態で) リクエストを送る
• service-worker は onfetch ハンドラーでリクエストに割り込み、main-thread に postMessage を送る
• main-thread の onmessage ハンドラーで、main-thread 上の操作を行い、その結果を serviceWorker に postMessage で返り返す
• service-worker の onmessage で結果を受け取り、worker の onfetch の response を返す

基本的に現代のフロントエンドでは同期 XHR は禁忌とされているが、 Web Worker 上では同期 xhr によってブロッキングされたとしても main-thread をブロックしない。ただし、この時 WebWorker のイベントループは止まっているので、worker thread 側の他の並列タスクは止まっている。

その上で、 クライアントオブジェクトへの参照は Proxy でラップする。Proxy オブジェクトでアクセスが発生するたびに、Proxy のリフレクションで同期XHRでメインスレッド側の値を解決する。

partytown の新規性

worker で同期 XHR という発想はなかった。partytown に限らず、この発想を応用すれば色々なことができそうなので、自分で実装してみる。

やってみた

この web worker 上で preact のコードが動くところまでやった。

/** @jsx h */
import { expose } from "comlink";
import { stubDocumentOnWorker } from "./town";
import { h, render } from "preact";
import { useEffect, useState } from "preact/hooks";

function App() {
  const [counter, setCounter] = useState(0);
  useEffect(() => {
    setInterval(() => {
      setCounter((n) => n + 1);
    }, 1000);
  }, []);
  return (
    <div>
      Hello,
      {counter}
    </div>
  );
}

const api = {
  start() {
    stubDocumentOnWorker();
    const el = document.createElement("div");
    render(<App />, el);
    document.body.appendChild(el);
  },
};

export type Api = typeof api;

expose(api);

この実装のサボっている点

色々サボっていて、イテレータを実装してないのと、メインスレッド側に関数ハンドラをセットする部分を実装してないので、onclick 等のイベントハンドラを扱えない。ここはやる気が出たらやる。

GitHub - immerjs/immer: Create the next immutable state by mutating the current one

実装解説

ソースコードはここで、要点を解説する

https://github.com/mizchi/mytown

ServiceWorker

まず service worker から実装した。onfetch でメインスレッドに postMessage を投げて、その結果を受けとる。

town.js

//...

let cnt = 0;
const _callbacks = new Map<number, (value: any) => void>();

export function handleMessageOnServiceWorker(ev: MessageEvent) {
  const payload = ev.data as ResponsePayload;
  const fn = _callbacks.get(payload.id);
  fn?.(payload.value);
  _callbacks.delete(ev.data.id);
}

export async function handleFetchOnServiceWorker(event: any) {
  const url = new URL(event.request.url);
  const encodedCmd = url.search.substr(1);
  const cmd = JSON.parse(atob(encodedCmd)) as RemoteCommand;
  const id = cnt++;
  let resolve: any;
  const promise = new Promise<any>((r) => (resolve = r));
  _callbacks.set(id, resolve);
  // @ts-ignore
  const client = await clients.get(event.clientId);
  client.postMessage({
    type: "req",
    cmd,
    id,
  } as RequestPayload);

  const value = await promise;
  return new Response(JSON.stringify(value));
}

unique な id を生成して callback の Map に Promise の resolve する関数を登録する。
serviceworker 側の onmessage で、その callback を解決して終了させる。

これを service-worker のリスナに登録する。リクエスト先を /__town?... のときだけ handleFetchOnServiceWorker で実装する。

src/sw.ts

import {
  handleFetchOnServiceWorker,
  handleMessageOnServiceWorker,
} from "./town";

const log = (...args: any) => console.log("[sw]", ...args);

self.addEventListener("install", (event: any) => {
  log("install", version);
  // @ts-ignore
  event.waitUntil(self.skipWaiting());
});

self.addEventListener("activate", (event: any) => {
  log("activate claim", version);
  // @ts-ignore
  event.waitUntil(self.clients.claim());
});

self.addEventListener("fetch", (event: any) => {
  const url = new URL(event.request.url);
  if (url.pathname.startsWith("/__town")) {
    log("handle", "/__town");
    return event.respondWith(handleFetchOnServiceWorker(event));
  }
});

self.addEventListener("message", handleMessageOnServiceWorker);

別にサーバーに飛ぶわけでもないし、 body のエンコーディングが面倒なので、/__town?<base64> で、 JSON を base64 エンコーディングを受け取ることにした。

worker の同期XHR

次に、worker からこのエンドポイントを同期XHRで呼ぶ関数を作る。

type Ptr = string | number;

type RemoteValue =
  | {
      isPtr: true;
      ptr: Ptr;
      parentPtr?: Ptr;
    }
  | {
      isPtr: false;
      value: any;
    };

export type RemoteCommand =
  | {
      op: "apply";
      ptr: Ptr;
      callerPtr?: Ptr;
      args: RemoteValue[];
    }
  | {
      op: "set";
      ptr: Ptr;
      key: string | number;
      value: RemoteValue;
    }
  | {
      op: "access";
      ptr: Ptr;
      key: string | number;
    };

export const execCommandSyncOnWorker = (cmd: RemoteCommand): RemoteValue => {
  const encoded = btoa(JSON.stringify(cmd));
  const url = "/__town?" + encoded;
  let result: any;
  const xhr = new XMLHttpRequest();
  xhr.open("GET", url, false);
  xhr.onload = (_e) =>
    xhr.readyState === 4 && xhr.status === 200 && (result = xhr.responseText);
  xhr.onerror = console.error;
  xhr.send(null);
  return JSON.parse(result);
};

メンバへのアクセス、代入、関数呼び出しのRPCを定義する。

この際、main-thread から worker に JSON シリアライズできないものを送ることができないので、その際はポインタIDを作って扱う。実際はこのIDを元に Proxy でラップして扱う。

main thread

ServiceWorker <=> Main の RPC の実行を実装する。

返り値のためにオブジェクトを生成するとき、worker には RemoteValue として適当に生成した id を返す。それは instanceMap に保存して id でインスタンスを見つけられるようにする。

// === main thread
const instanceMap = new Map<Ptr, any>();
if (globalThis.document) {
  instanceMap.set("document", document);
  instanceMap.set("window", window);
}

function isTransferrable(raw: any) {
  return raw == null || ["number", "string", "boolean"].includes(typeof raw);
}

const raw2remote = (raw: any): RemoteValue => {
  if (isTransferrable(raw)) {
    return {
      isPtr: false,
      value: raw,
    };
  } else {
    const newPtr = Math.random().toString(32).substr(2);
    instanceMap.set(newPtr, raw);
    return {
      isPtr: true,
      ptr: newPtr,
    };
  }
};
export function resolveCommandOnMain(command: RemoteCommand): RemoteValue {
  switch (command.op) {
    case "access": {
      const parent = instanceMap.get(command.ptr);
      const rawValue = parent[command.key];
      return raw2remote(rawValue);
    }
    case "set": {
      const parent = instanceMap.get(command.ptr);
      const rightValue = command.value.isPtr
        ? instanceMap.get(command.value.ptr)
        : command.value.value;
      parent[command.key] = rightValue;
      return raw2remote(undefined);
    }
    case "apply": {
      const fn = instanceMap.get(command.ptr);
      const rawArgs = command.args.map((value) => {
        return value.isPtr ? instanceMap.get(value.ptr) : value.value;
      });
      let rawValue;
      if (command.callerPtr) {
        const caller = instanceMap.get(command.callerPtr);
        rawValue = fn.apply(caller, rawArgs);
      } else {
        rawValue = fn(...rawArgs);
      }
      return raw2remote(rawValue);
    }
  }
}

apply の関数呼び出しに注意が必要で、関数参照を直接呼び出すと this が解決できないことがあり、 document.querySelector(...) みたいなパターンのために呼び出し元の ptr を一緒に渡す。

これを postMessage の onmessage ハンドラとして渡す

// こういう型を定義してある
export type RequestPayload = {
  type: "req";
  id: number;
  cmd: RemoteCommand;
};

export type ResponsePayload = {
  type: "res";
  id: number;
  value: any;
};

export function handleMessageOnMain(event: MessageEvent) {
  if (event.data.type === "req") {
    const req = event.data as RequestPayload;
    navigator.serviceWorker.controller!.postMessage({
      type: "res",
      id: req.id,
      value: resolveCommandOnMain(req.cmd),
    } as ResponsePayload);
  }
}

Worker の Mock オブジェクト

Proxy を使ってオブジェクトに対する get/set/apply を実装する。順番が前後したが Command RPCの命名はここから決めている。

apply でエラーにならないように適当な関数を元に Proxy を作る。

アクセスパターンに応じて、先に準備した

export function createPtr(ptr: Ptr, parentPtr?: Ptr): any {
  return new Proxy(() => {}, {
    apply(_target, _thisArg, argumentsList) {
      const ret = execCommandSyncOnWorker({
        op: "apply",
        ptr: ptr,
        callerPtr: parentPtr,
        args: argumentsList.map((arg) => {
          if (isTransferrable(arg)) {
            return {
              isPtr: false,
              value: arg,
            };
          }
          return (
            arg._ptr ?? {
              isPtr: false,
              value: arg,
            }
          );
        }),
      } as RemoteCommand);

      if (ret.isPtr) {
        return createPtr(ret.ptr);
      } else {
        return ret.value;
      }
    },
    set(_target, propertyName, value, _receiver) {
      const remoteValue =
        typeof value === "object" && value._ptr
          ? value._ptr
          : {
              isPtr: false,
              value,
            };
      const _ret = execCommandSyncOnWorker({
        op: "set",
        ptr: ptr,
        key: propertyName,
        value: remoteValue,
      } as RemoteCommand);
      return true;
    },
    get(_target, propertyName) {
      if (propertyName == "_ptr") {
        return {
          ptr,
          isPtr: true,
        };
      }
      const ret = execCommandSyncOnWorker({
        op: "access",
        ptr: ptr,
        key: propertyName,
      } as RemoteCommand);

      if (ret.isPtr) {
        return createPtr(ret.ptr, ptr);
      } else {
        return ret.value;
      }
    },
  });
}

export const stubDocumentOnWorker = () => {
  globalThis.document = createPtr("document");
  globalThis.window = createPtr("window");
};

createPtr() は再帰的な構造になっていて、これで生成されたオブジェクトはメンバーアクセスで先に用意した同期XHRを呼び、メインスレッドから解決された RemoteValue を createPtr() でラップして返す。

これを worker 環境では globalThis.document に document という名前で登録する。(window も)

main-theard のところで解説しなかったが、main-theard の instanceMap でこの2つを解決できるようにインスタンスを登録してある。

// main thread
instanceMap.set("document", document);
instanceMap.set("window", window);

これで window と document はリモートオブジェクトとして扱えるようになる。

実装したDOM を worker から使う

worker を起動するところなどは割愛するが、(vite + comlink でやってる) これで偽の document オブジェクトを worker から操作すると、メインスレッドで同じ操作が適用される。

というのを実証するために、 preact を動かしてみた。

/** @jsx h */

import { expose } from "comlink";
import { stubDocumentOnWorker } from "./town";
import { h, render } from "preact";
import { useEffect, useState } from "preact/hooks";

function App() {
  const [counter, setCounter] = useState(0);
  useEffect(() => {
    setInterval(() => {
      setCounter((n) => n + 1);
    }, 1000);
  }, []);
  return (
    <div>
      Hello,
      {counter}
    </div>
  );
}

const api = {
  start() {
    stubDocumentOnWorker();
    const el = document.createElement("div");
    render(<App />, el);
    document.body.appendChild(el);
  },
};

export type Api = typeof api;

expose(api);

この Worker を comlink 経由で start を叩く実装をする。

最初に service-worker を登録する。

// main-thread
import Worker from "./worker?worker";
import { wrap, Remote } from "comlink";
import type { Api } from "./worker";
import { handleMessageOnMain } from "./town";

const api = wrap(new Worker()) as Remote<Api>;

async function main() {
  const _reg = await navigator.serviceWorker.register("/sw.js");
  navigator.serviceWorker.addEventListener("message", handleMessageOnMain);
  const postButton = document.createElement("button");
  postButton.onclick = () => {
    api.start();
  };
  postButton.textContent = "start";
  document.body.appendChild(postButton);
}

main();

というわけでこれが動く

おわり

Proxy オブジェクトのテクニックは自力で実装したがエッジケースが無限にあるので、ユースケースを絞って実装したほうがよさそう。

プロキシオブジェクトへのメンバへのアクセスの度にプロセスまたぎの非同期アクセスが走りまくるので、あまりパフォーマンスはでなさそう。そう考えると結局副作用をバッチできる worker-dom と同じ実装になっていく気がする。

見かけ上同期に見せかける方法は便利そうなので、いずれなにかに使えそう。