feat(electron): expose electron apis to web worker (#9441)

fix AF-2044

packages/frontend/apps/electron/src/preload/bootstrap.ts

@@ -2,12 +2,11 @@ import '@sentry/electron/preload';
 
 import { contextBridge } from 'electron';
 
-import { appInfo, getElectronAPIs } from './electron-api';
+import { apis, appInfo, events, requestWebWorkerPort } from './electron-api';
 import { sharedStorage } from './shared-storage';
 
-const { apis, events } = getElectronAPIs();
-
 contextBridge.exposeInMainWorld('__appInfo', appInfo);
 contextBridge.exposeInMainWorld('__apis', apis);
 contextBridge.exposeInMainWorld('__events', events);
 contextBridge.exposeInMainWorld('__sharedStorage', sharedStorage);
+contextBridge.exposeInMainWorld('__requestWebWorkerPort', requestWebWorkerPort);

packages/frontend/apps/electron/src/preload/electron-api.ts

@@ -13,22 +13,6 @@ import {
   type RendererToHelper,
 } from '../shared/type';
 
-export function getElectronAPIs() {
-  const mainAPIs = getMainAPIs();
-  const helperAPIs = getHelperAPIs();
-
-  return {
-    apis: {
-      ...mainAPIs.apis,
-      ...helperAPIs.apis,
-    },
-    events: {
-      ...mainAPIs.events,
-      ...helperAPIs.events,
-    },
-  };
-}
-
 type Schema =
   | 'affine'
   | 'affine-canary'
@@ -248,3 +232,60 @@ function getHelperAPIs() {
     return { apis: {}, events: {} };
   }
 }
+
+const mainAPIs = getMainAPIs();
+const helperAPIs = getHelperAPIs();
+
+export const apis = {
+  ...mainAPIs.apis,
+  ...helperAPIs.apis,
+};
+
+export const events = {
+  ...mainAPIs.events,
+  ...helperAPIs.events,
+};
+
+// Create MessagePort that can be used by web workers
+export function requestWebWorkerPort() {
+  const ch = new MessageChannel();
+
+  const localPort = ch.port1;
+  const remotePort = ch.port2;
+
+  // todo: should be able to let the web worker use the electron APIs directly for better performance
+  const flattenedAPIs = Object.entries(apis).flatMap(([namespace, api]) => {
+    return Object.entries(api as any).map(([method, fn]) => [
+      `${namespace}:${method}`,
+      fn,
+    ]);
+  });
+
+  AsyncCall(Object.fromEntries(flattenedAPIs), {
+    channel: createMessagePortChannel(localPort),
+    log: false,
+  });
+
+  const cleanup = () => {
+    remotePort.close();
+    localPort.close();
+  };
+
+  const portId = crypto.randomUUID();
+
+  setTimeout(() => {
+    window.postMessage(
+      {
+        type: 'electron:request-api-port',
+        portId,
+        ports: [remotePort],
+      },
+      '*',
+      [remotePort]
+    );
+  });
+
+  localPort.start();
+
+  return { portId, cleanup };
+}

packages/frontend/core/src/modules/docs-search/worker/in-worker.ts

@@ -1,3 +1,4 @@
+import { getElectronAPIs } from '@affine/electron-api/web-worker';
 import type {
   AttachmentBlockModel,
   BookmarkBlockModel,
@@ -43,6 +44,11 @@ const LRU_CACHE_SIZE = 5;
 // lru cache for ydoc instances, last used at the end of the array
 const lruCache = [] as { doc: YDoc; hash: string }[];
 
+const electronAPIs = BUILD_CONFIG.isElectron ? getElectronAPIs() : null;
+
+// @ts-expect-error test
+globalThis.__electronAPIs = electronAPIs;
+
 async function digest(data: Uint8Array) {
   if (
     globalThis.crypto &&

packages/frontend/core/src/modules/docs-search/worker/out-worker.ts

@@ -1,4 +1,5 @@
 import { DebugLogger } from '@affine/debug';
+import { connectWebWorker } from '@affine/electron-api/web-worker';
 import { MANUALLY_STOP, throwIfAborted } from '@toeverything/infra';
 
 import type {
@@ -12,6 +13,7 @@ const logger = new DebugLogger('affine:indexer-worker');
 
 export async function createWorker(abort: AbortSignal) {
   let worker: Worker | null = null;
+  let electronApiCleanup: (() => void) | null = null;
   while (throwIfAborted(abort)) {
     try {
       worker = await new Promise<Worker>((resolve, reject) => {
@@ -29,6 +31,11 @@ export async function createWorker(abort: AbortSignal) {
           }
         });
         worker.postMessage({ type: 'init', msgId: 0 } as WorkerIngoingMessage);
+
+        if (BUILD_CONFIG.isElectron) {
+          electronApiCleanup = connectWebWorker(worker);
+        }
+
         setTimeout(() => {
           reject('timeout');
         }, 1000 * 30 /* 30 sec */);
@@ -97,6 +104,7 @@ export async function createWorker(abort: AbortSignal) {
     dispose: () => {
       terminateAbort.abort(MANUALLY_STOP);
       worker.terminate();
+      electronApiCleanup?.();
     },
   };
 }

packages/frontend/core/src/modules/workbench/view/split-view/panel.tsx

@@ -9,12 +9,12 @@ import {
 import { useSortable } from '@dnd-kit/sortable';
 import { useLiveData, useService } from '@toeverything/infra';
 import { assignInlineVars } from '@vanilla-extract/dynamic';
-import type { SetStateAction } from 'jotai';
 import type {
   Dispatch,
   HTMLAttributes,
   PropsWithChildren,
   RefObject,
+  SetStateAction,
 } from 'react';
 import {
   memo,

packages/frontend/electron-api/package.json

@@ -5,6 +5,10 @@
   "private": true,
   "main": "./src/index.ts",
   "exports": {
-    ".": "./src/index.ts"
+    ".": "./src/index.ts",
+    "./web-worker": "./src/web-worker.ts"
+  },
+  "dependencies": {
+    "async-call-rpc": "^6.4.2"
   }
-}
+}

packages/frontend/electron-api/src/web-worker.ts

@@ -0,0 +1,83 @@
+import { AsyncCall, type EventBasedChannel } from 'async-call-rpc';
+
+import type { ClientHandler } from '.';
+
+const WORKER_PORT_MESSAGE_TYPE = 'electron-api-port';
+
+// connect web worker to preload, so that the web worker can use the electron APIs
+export function connectWebWorker(worker: Worker) {
+  const { portId, cleanup } = (globalThis as any).__requestWebWorkerPort();
+
+  const portMessageListener = (event: MessageEvent) => {
+    if (
+      event.data.type === 'electron:request-api-port' &&
+      event.data.portId === portId
+    ) {
+      const [port] = event.data.ports as MessagePort[];
+
+      // worker should be ready to receive message
+      worker.postMessage(
+        {
+          type: WORKER_PORT_MESSAGE_TYPE,
+          ports: [port],
+        },
+        [port]
+      );
+    }
+  };
+
+  window.addEventListener('message', portMessageListener);
+
+  return () => {
+    window.removeEventListener('message', portMessageListener);
+    cleanup();
+  };
+}
+
+const createMessagePortChannel = (port: MessagePort): EventBasedChannel => {
+  return {
+    on(listener) {
+      port.onmessage = e => {
+        listener(e.data);
+      };
+      port.start();
+      return () => {
+        port.onmessage = null;
+        try {
+          port.close();
+        } catch (err) {
+          console.error('[worker] close port error', err);
+        }
+      };
+    },
+    send(data) {
+      port.postMessage(data);
+    },
+  };
+};
+
+// get the electron APIs for the web worker (should be called in the web worker)
+export function getElectronAPIs(): ClientHandler {
+  const { promise, resolve } = Promise.withResolvers<MessagePort>();
+  globalThis.addEventListener('message', event => {
+    if (event.data.type === WORKER_PORT_MESSAGE_TYPE) {
+      const [port] = event.ports;
+      resolve(port);
+    }
+  });
+
+  const rpc = AsyncCall<Record<string, any>>(null, {
+    channel: promise.then(p => createMessagePortChannel(p)),
+    log: false,
+  });
+
+  return new Proxy<ClientHandler>(rpc as any, {
+    get(_, namespace: string) {
+      return new Proxy(rpc as any, {
+        get(_, method: string) {
+          return rpc[`${namespace}:${method}`];
+        },
+      });
+    },
+  });
+}

yarn.lock

@@ -454,6 +454,8 @@ __metadata:
 "@affine/electron-api@workspace:*, @affine/electron-api@workspace:packages/frontend/electron-api":
   version: 0.0.0-use.local
   resolution: "@affine/electron-api@workspace:packages/frontend/electron-api"
+  dependencies:
+    async-call-rpc: "npm:^6.4.2"
   languageName: unknown
   linkType: soft