feat: client-worker communication

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

@@ -17,20 +17,26 @@ export default function SharedWorker(options?: { root?: string }): Plugin {
       if (id.startsWith(root)) {
         const isClient = !id.endsWith('worker_file');
 
-        const exports = findExports(code).filter((ex) => !!ex.name);
+        const exports = findExports(code)
+          .filter((ex) => !!ex.name)
+          .filter((ex) => ex.name !== 'dispatch' && ex.name !== 'addMessageListener');
 
         if (isClient) {
           const workerName = path.basename(id.split('?')[0]);
           return [
             `// vite-plugin-sharedworker starts`,
-            `import { defineClient } from 'vite-plugin-sharedworker/runtime'`,
-            `const worker = new SharedWorker(new URL(${JSON.stringify(
+            `import { defineClientFactory } from 'vite-plugin-sharedworker/runtime'`,
+            `const __worker__ = new SharedWorker(new URL(${JSON.stringify(
               id
             )}, import.meta.url), { type: 'module', name: ${JSON.stringify(workerName)} });`,
-            `const client = defineClient(worker)`,
+            `const __factory__ = defineClientFactory(__worker__)`,
+            `const client = __factory__.defineClient()`,
+            `export const dispatch = client.dispatch`,
+            `export const addMessageListener = client.addMessageListener`,
             `// vite-plugin-sharedworker ends`,
             ...exports.map(
-              (ex) => `export const ${ex.name} = client.defineFunction(${JSON.stringify(ex.name)})`
+              (ex) =>
+                `export const ${ex.name} = __factory__.defineFunction(${JSON.stringify(ex.name)})`
             )
           ].join('\n');
         } else {
@@ -40,6 +46,7 @@ export default function SharedWorker(options?: { root?: string }): Plugin {
             `const worker = defineSharedWorker(self, [${exports
               .map((ex) => ex.name!)
               .join(', ')}])`,
+            `const client = { dispatch: undefined, addMessageListener: undefined }`,
             `// vite-plugin-sharedworker ends`,
             ''
           ];

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

@@ -1,7 +1,7 @@
 /// <reference lib="DOM" />
 /// <reference lib="WebWorker" />
 
-import type { SharedWorkerServer } from './types';
+import type { SharedWorkerServer, SharedWorkerClient, MessageCallback } from './types';
 
 import {
   makeRpcPayload,
@@ -18,9 +18,11 @@ export * from './types';
 export * from './payload';
 
 type _SharedWorkerServer = SharedWorkerServer;
+type _SharedWorkerClient = SharedWorkerClient;
 
 declare global {
   export type SharedWorkerServer = _SharedWorkerServer;
+  export type SharedWorkerClient = _SharedWorkerClient;
 }
 
 type Fn = (...args: any[]) => any;
@@ -45,6 +47,8 @@ export function defineSharedWorker(self: SharedWorkerGlobalScope, fns: Fn[]): Sh
     }
   }, 1000);
 
+  const messageCallbacks: MessageCallback<any>[] = [];
+
   self.addEventListener('connect', (event) => {
     const port = event.ports[0];
     ports.set(port, new Date());
@@ -60,6 +64,8 @@ export function defineSharedWorker(self: SharedWorkerGlobalScope, fns: Fn[]): Sh
           } else {
             console.error(`Unknown message: ${JSON.stringify(payload, null, 2)}`);
           }
+        } else if (payload.command === 'broadcast') {
+          await Promise.all(messageCallbacks.map((fn) => fn.apply(event, [payload.data])));
         } else if (payload.command === 'ping') {
           ports.set(port, new Date());
         }
@@ -76,6 +82,9 @@ export function defineSharedWorker(self: SharedWorkerGlobalScope, fns: Fn[]): Sh
     dispatch(port, data: any) {
       port.postMessage(makeBroadcastPayload(data));
     },
+    addMessageListener(fn) {
+      messageCallbacks.push(fn);
+    },
     broadcast(data: any) {
       for (const port of ports.keys()) {
         port.postMessage(makeBroadcastPayload(data));
@@ -84,12 +93,13 @@ export function defineSharedWorker(self: SharedWorkerGlobalScope, fns: Fn[]): Sh
   };
 }
 
-export function defineClient(worker: SharedWorker) {
+export function defineClientFactory(worker: SharedWorker) {
   worker.port.start();
 
   const callbacks = new Map<string, (payload: RpcPayload) => void>();
+  const messageCallbacks: MessageCallback<any>[] = [];
 
-  worker.port.addEventListener('message', (event) => {
+  worker.port.addEventListener('message', async (event) => {
     const payload = parsePayload(event.data);
     if (payload) {
       if (payload.command === 'rpc') {
@@ -101,6 +111,7 @@ export function defineClient(worker: SharedWorker) {
           console.error(`Unknown message: ${JSON.stringify(payload, null, 2)}`);
         }
       } else if (payload.command === 'broadcast') {
+        await Promise.all(messageCallbacks.map((fn) => fn.apply(event, [payload.data])));
       }
     }
   });
@@ -122,6 +133,16 @@ export function defineClient(worker: SharedWorker) {
           worker.port.postMessage(payload);
         });
       };
+    },
+    defineClient(): SharedWorkerClient {
+      return {
+        addMessageListener(fn) {
+          messageCallbacks.push(fn);
+        },
+        dispatch(data) {
+          worker.port.postMessage(makeBroadcastPayload(data));
+        }
+      };
     }
   };
 }

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

@@ -38,6 +38,8 @@ export function parsePayload(payload: Payload): Payload | undefined {
     return payload;
   } else if (payload.command === 'rpc') {
     return payload;
+  } else if (payload.command === 'broadcast') {
+    return payload;
   } else {
     return undefined;
   }

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

@@ -1,7 +1,17 @@
+export type MessageCallback<T> = (this: Event, data: T) => void | Promise<void>;
+
 export interface SharedWorkerServer {
   ports(): MessagePort[];
 
+  addMessageListener<T>(fn: MessageCallback<T>): void;
+
   dispatch(port: MessagePort, data: any): void;
 
   broadcast(data: any): void;
 }
+
+export interface SharedWorkerClient {
+  addMessageListener<T>(fn: MessageCallback<T>): void;
+
+  dispatch(data: any): void;
+}

playground/src/main.ts

@@ -1,7 +1,11 @@
 import 'uno.css';
 import '@onekuma/reset/tailwind.css';
 
-import { add, sub } from '../worker';
+import { add, sub, addMessageListener } from '../worker';
+
+addMessageListener((data) => {
+  console.log(data);
+});
 
 async function bootstrap() {
   console.log(await add(1, 2));

playground/worker/index.ts

@@ -1,5 +1,16 @@
 declare const worker: SharedWorkerServer;
 
+declare const client: SharedWorkerClient;
+
+export const dispatch = client.dispatch;
+
+export const addMessageListener = client.addMessageListener;
+
+setInterval(() => {
+  console.log('Hello');
+  worker.broadcast('Hello, this is sharedworker');
+}, 2000);
+
 export async function add(a: number, b: number) {
   console.log(worker.ports());
   return a + b;