From 887732179edb35f49bcbaf4eafe5edb7a2be938c Mon Sep 17 00:00:00 2001 From: pengx17 Date: Tue, 31 Dec 2024 03:17:02 +0000 Subject: [PATCH] feat(electron): expose electron apis to web worker (#9441) fix AF-2044 --- .../apps/electron/src/preload/bootstrap.ts | 5 +- .../apps/electron/src/preload/electron-api.ts | 73 ++++++++++++---- .../modules/docs-search/worker/in-worker.ts | 6 ++ .../modules/docs-search/worker/out-worker.ts | 8 ++ .../workbench/view/split-view/panel.tsx | 2 +- packages/frontend/electron-api/package.json | 8 +- .../frontend/electron-api/src/web-worker.ts | 83 +++++++++++++++++++ yarn.lock | 2 + 8 files changed, 165 insertions(+), 22 deletions(-) create mode 100644 packages/frontend/electron-api/src/web-worker.ts diff --git a/packages/frontend/apps/electron/src/preload/bootstrap.ts b/packages/frontend/apps/electron/src/preload/bootstrap.ts index be1ebaf3db089..8bbce97528b54 100644 --- a/packages/frontend/apps/electron/src/preload/bootstrap.ts +++ b/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); diff --git a/packages/frontend/apps/electron/src/preload/electron-api.ts b/packages/frontend/apps/electron/src/preload/electron-api.ts index 3f0d897e5bfcd..5706a3d7857c8 100644 --- a/packages/frontend/apps/electron/src/preload/electron-api.ts +++ b/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 }; +} diff --git a/packages/frontend/core/src/modules/docs-search/worker/in-worker.ts b/packages/frontend/core/src/modules/docs-search/worker/in-worker.ts index 8a0e68360b249..ceb33c408e38a 100644 --- a/packages/frontend/core/src/modules/docs-search/worker/in-worker.ts +++ b/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 && diff --git a/packages/frontend/core/src/modules/docs-search/worker/out-worker.ts b/packages/frontend/core/src/modules/docs-search/worker/out-worker.ts index 4c2de87219bcd..97e45dc2bd7a4 100644 --- a/packages/frontend/core/src/modules/docs-search/worker/out-worker.ts +++ b/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((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?.(); }, }; } diff --git a/packages/frontend/core/src/modules/workbench/view/split-view/panel.tsx b/packages/frontend/core/src/modules/workbench/view/split-view/panel.tsx index 8a8390f015ee8..65f3aaeab5d79 100644 --- a/packages/frontend/core/src/modules/workbench/view/split-view/panel.tsx +++ b/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, diff --git a/packages/frontend/electron-api/package.json b/packages/frontend/electron-api/package.json index 67907543f7c50..7fe7666e70f9c 100644 --- a/packages/frontend/electron-api/package.json +++ b/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" } -} +} \ No newline at end of file diff --git a/packages/frontend/electron-api/src/web-worker.ts b/packages/frontend/electron-api/src/web-worker.ts new file mode 100644 index 0000000000000..ed986e46f5eec --- /dev/null +++ b/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(); + globalThis.addEventListener('message', event => { + if (event.data.type === WORKER_PORT_MESSAGE_TYPE) { + const [port] = event.ports; + resolve(port); + } + }); + + const rpc = AsyncCall>(null, { + channel: promise.then(p => createMessagePortChannel(p)), + log: false, + }); + + return new Proxy(rpc as any, { + get(_, namespace: string) { + return new Proxy(rpc as any, { + get(_, method: string) { + return rpc[`${namespace}:${method}`]; + }, + }); + }, + }); +} diff --git a/yarn.lock b/yarn.lock index 2efa2bca563cc..bb8d7c0897772 100644 --- a/yarn.lock +++ b/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