feat: shared worker rpc

packages/vite-plugin-worker/src/index.ts

@@ -17,24 +17,26 @@ export default function SharedWorker(options?: { root?: string }): Plugin {
       if (id.startsWith(root)) {
         const isClient = !id.endsWith('worker_file');
 
-        // console.log(id);
-        if (isClient) {
-          const exports = findExports(code);
+        const exports = findExports(code).filter((ex) => !!ex.name);
 
+        if (isClient) {
           return [
+            `// vite-plugin-sharedworker starts`,
+            `import { defineClient } from 'vite-plugin-sharedworker/runtime'`,
             `const worker = new SharedWorker(new URL(${JSON.stringify(
               id
             )}, import.meta.url), { type: 'module', name: 'worker' });`,
-            `worker.port.start()`,
-            ...exports
-              .filter((ex) => ex.name)
-              .map((ex) => `export const ${ex.name} = (...args: any[]) => {}`)
+            `const client = defineClient(worker)`,
+            `// vite-plugin-sharedworker ends`,
+            ...exports.map(
+              (ex) => `export const ${ex.name} = client.defineFunction(${JSON.stringify(ex.name)})`
+            )
           ].join('\n');
         } else {
           const imports = [
             `// vite-plugin-sharedworker starts`,
             `import { defineSharedWorker } from 'vite-plugin-sharedworker/runtime'`,
-            `defineSharedWorker(self)`,
+            `defineSharedWorker(self, [${exports.map((ex) => ex.name!).join(', ')}])`,
             `// vite-plugin-sharedworker ends`,
             ''
           ];

packages/vite-plugin-worker/src/runtime.ts

@@ -1,12 +1,85 @@
+/// <reference lib="DOM" />
 /// <reference lib="WebWorker" />
 
-export function defineSharedWorker(self: SharedWorkerGlobalScope) {
+import { randomString } from './utils';
+
+type Fn = (...args: any[]) => any;
+
+export function defineSharedWorker(self: SharedWorkerGlobalScope, fns: Fn[]) {
+  const map = new Map<string, Fn>();
+  for (const fn of fns) {
+    map.set(fn.name, fn);
+  }
+
   self.addEventListener('connect', (event) => {
-    console.log('connect', event);
     const port = event.ports[0];
-    port.addEventListener('message', (event) => {
-      console.log('message', event);
+
+    port.addEventListener('message', async (event) => {
+      const payload = parsePayload(event.data);
+      const fn = map.get(payload.name);
+      if (fn) {
+        const result = await fn.apply(event, payload.arguments);
+        port.postMessage(makePayload(payload.id, payload.name, result));
+      } else {
+        console.error(`Unknown message: ${JSON.stringify(payload, null, 2)}`);
+      }
     });
+
     port.start();
   });
 }
+
+export function defineClient(worker: SharedWorker) {
+  worker.port.start();
+
+  const callbacks = new Map<string, (payload: Payload) => void>();
+
+  worker.port.addEventListener('message', (event) => {
+    const payload = parsePayload(event.data);
+    const callback = callbacks.get(payload.id);
+    if (callback) {
+      callback(payload);
+    } else {
+      console.error(`Unknown message: ${JSON.stringify(payload, null, 2)}`);
+    }
+  });
+
+  return {
+    defineFunction(name: string) {
+      return (...args: any[]) => {
+        const id = randomString();
+        const payload = makePayload(id, name, args);
+
+        return new Promise((res) => {
+          callbacks.set(id, (payload) => {
+            res(payload.arguments);
+          });
+          worker.port.postMessage(payload);
+        });
+      };
+    }
+  };
+}
+
+interface Payload {
+  id: string;
+  name: string;
+  arguments: any;
+}
+
+function makePayload(id: string, name: string, args: any) {
+  return JSON.stringify({
+    i: id,
+    f: name,
+    a: args
+  });
+}
+
+function parsePayload(payload: string) {
+  const p = JSON.parse(payload);
+  return {
+    id: p.i,
+    name: p.f as string,
+    arguments: p.a as any[]
+  };
+}

packages/vite-plugin-worker/src/utils.ts

@@ -0,0 +1,10 @@
+function random(l: number, r: number): number {
+  return l + Math.round(Math.random() * (r - l));
+}
+
+const character_table = '0123456789abcdefghijklmnopqrstuvwxyz';
+export function randomString(length = 8): string {
+  return Array.apply(null, Array(length))
+    .map(() => character_table[random(0, character_table.length - 1)])
+    .join('');
+}

playground/main.ts

@@ -1,16 +1,10 @@
 import { add } from './worker';
 
-console.log(add(1, 2));
-
-const worker = new SharedWorker(new URL('./worker/index.ts', import.meta.url), {
-  type: 'module',
-  name: 'worker'
-});
-
-worker.port.start();
-
-worker.port.postMessage('hello 1');
-
-console.log('OK');
-
-export {};
+console.log(await add(1, 2));
+
+// const worker = new SharedWorker(new URL('./worker/index.ts', import.meta.url), {
+//   type: 'module',
+//   name: 'worker'
+// });
+// worker.port.start();
+// worker.port.postMessage('hello 1');

playground/worker/index.ts

@@ -4,6 +4,6 @@
 // defineSharedWorker(self);
 // export type {};
 
-export function add(a: number, b: number) {
+export async function add(a: number, b: number) {
   return a + b;
 }