feat: check port alive

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

@@ -37,7 +37,9 @@ export default function SharedWorker(options?: { root?: string }): Plugin {
           const imports = [
             `// vite-plugin-sharedworker starts`,
             `import { defineSharedWorker } from 'vite-plugin-sharedworker/runtime'`,
-            `defineSharedWorker(self, [${exports.map((ex) => ex.name!).join(', ')}])`,
+            `const worker = defineSharedWorker(self, [${exports
+              .map((ex) => ex.name!)
+              .join(', ')}])`,
             `// vite-plugin-sharedworker ends`,
             ''
           ];

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

@@ -0,0 +1,117 @@
+/// <reference lib="DOM" />
+/// <reference lib="WebWorker" />
+
+import type { SharedWorkerServer } from './types';
+
+import {
+  makeRpcPayload,
+  parsePayload,
+  RpcPayload,
+  makePingPayload,
+  makeBroadcastPayload
+} from './payload';
+
+import { randomString } from './utils';
+
+export * from './types';
+
+type _SharedWorkerServer = SharedWorkerServer;
+
+declare global {
+  export type SharedWorkerServer = _SharedWorkerServer;
+}
+
+type Fn = (...args: any[]) => any;
+
+export function defineSharedWorker(self: SharedWorkerGlobalScope, fns: Fn[]): SharedWorkerServer {
+  const map = new Map<string, Fn>();
+  for (const fn of fns) {
+    map.set(fn.name, fn);
+  }
+
+  const ports = new Map<MessagePort, Date>();
+
+  setInterval(() => {
+    const now = new Date().getTime();
+    for (const [port, date] of ports) {
+      if (now - date.getTime() >= 2000) {
+        ports.delete(port);
+      }
+    }
+  }, 1000);
+
+  self.addEventListener('connect', (event) => {
+    const port = event.ports[0];
+    ports.set(port, new Date());
+
+    port.addEventListener('message', async (event) => {
+      const payload = parsePayload(event.data);
+      if (payload) {
+        if (payload.command === 'rpc') {
+          const fn = map.get(payload.data.name);
+          if (fn) {
+            const result = await fn.apply(event, payload.data.args);
+            port.postMessage(makeRpcPayload(payload.data.id, payload.data.name, result));
+          } else {
+            console.error(`Unknown message: ${JSON.stringify(payload, null, 2)}`);
+          }
+        } else if (payload.command === 'ping') {
+          ports.set(port, new Date());
+        }
+      }
+    });
+
+    port.start();
+  });
+
+  return {
+    ports() {
+      return [...ports.keys()];
+    },
+    broadcast(data: any) {
+      for (const port of ports.keys()) {
+        port.postMessage(makeBroadcastPayload(data));
+      }
+    }
+  };
+}
+
+export function defineClient(worker: SharedWorker) {
+  worker.port.start();
+
+  const callbacks = new Map<string, (payload: RpcPayload) => void>();
+
+  worker.port.addEventListener('message', (event) => {
+    const payload = parsePayload(event.data);
+    if (payload) {
+      if (payload.command === 'rpc') {
+        const callback = callbacks.get(payload.data.id);
+        if (callback) {
+          callback(payload.data);
+        } else {
+          console.error(`Unknown message: ${JSON.stringify(payload, null, 2)}`);
+        }
+      }
+    }
+  });
+
+  setInterval(() => {
+    worker.port.postMessage(makePingPayload());
+  }, 500);
+
+  return {
+    defineFunction(name: string) {
+      return (...args: any[]) => {
+        const id = randomString();
+        const payload = makeRpcPayload(id, name, args);
+
+        return new Promise((res) => {
+          callbacks.set(id, (payload) => {
+            res(payload.args);
+          });
+          worker.port.postMessage(payload);
+        });
+      };
+    }
+  };
+}

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

@@ -0,0 +1,44 @@
+export interface RpcPayload {
+  id: string;
+  name: string;
+  args: any;
+}
+
+export type Payload =
+  | { command: 'rpc'; data: RpcPayload }
+  | { command: 'broadcast'; data: any }
+  | { command: 'ping' };
+
+export function makeRpcPayload(id: string, name: string, args: any): Payload {
+  return {
+    command: 'rpc',
+    data: {
+      id,
+      name,
+      args: args
+    }
+  };
+}
+
+export function makeBroadcastPayload(data: any): Payload {
+  return {
+    command: 'broadcast',
+    data
+  };
+}
+
+export function makePingPayload(): Payload {
+  return {
+    command: 'ping'
+  };
+}
+
+export function parsePayload(payload: Payload): Payload | undefined {
+  if (payload.command === 'ping') {
+    return payload;
+  } else if (payload.command === 'rpc') {
+    return payload;
+  } else {
+    return undefined;
+  }
+}

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

@@ -0,0 +1,5 @@
+export interface SharedWorkerServer {
+  ports(): MessagePort[];
+
+  broadcast(data: any): void;
+}

playground/tsconfig.json

@@ -15,6 +15,8 @@
     "strict": true,
     "target": "ESNext",
     "types": [
+      "vite/client",
+      "vite-plugin-sharedworker/runtime",
       "vite-plugin-info/client"
     ],
     "useDefineForClassFields": true

playground/tsconfig.worker.json

@@ -1,8 +1,15 @@
 {
   "compilerOptions": {
-    "lib": ["ESNext", "WebWorker"],
-    "types": ["vite/client"]
+    "lib": [
+      "ESNext",
+      "WebWorker"
+    ],
+    "types": [
+      "vite-plugin-sharedworker/runtime"
+    ]
   },
-  "include": ["./worker/*.ts"],
+  "include": [
+    "./worker/*.ts"
+  ],
   "extends": "./tsconfig.json"
 }

playground/worker/index.ts

@@ -1,13 +1,11 @@
-// /// <reference lib="WebWorker" />
-// declare const self: SharedWorkerGlobalScope;
-// import { defineSharedWorker } from 'vite-plugin-sharedworker/runtime';
-// defineSharedWorker(self);
-// export type {};
+declare const worker: SharedWorkerServer;
 
 export async function add(a: number, b: number) {
+  console.log(worker.ports());
   return a + b;
 }
 
 export async function sub(a: number, b: number) {
+  console.log(worker.ports());
   return a - b;
 }