diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts index 2dd68e2ffb..26c6ad9096 100644 --- a/desktop/src/main/services/ml-util-test.ts +++ b/desktop/src/main/services/ml-util-test.ts @@ -12,12 +12,18 @@ process.parentPort.once("message", (e) => { }); }); -/** Our hand-rolled IPC handler and router */ +/** + * Our hand-rolled IPC handler and router - the Node.js utility process end. + * + * Sibling of the electronMLWorker function (in `ml/worker.ts`) in the web code. + */ const handleMessage = async (m: unknown) => { - if (m && typeof m == "object" && "type" in m) { + if (m && typeof m == "object" && "type" in m && "id" in m) { + const id = m.id; switch (m.type) { case "foo": - if ("a" in m && typeof m.a == "string") return await foo(m.a); + if ("a" in m && typeof m.a == "string") + return { id, data: await foo(m.a) }; break; } } diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 043c339e3d..5e00f47e88 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -9,6 +9,7 @@ import { ensure } from "@/utils/ensure"; import { wait } from "@/utils/promise"; import { DOMParser } from "@xmldom/xmldom"; import { expose } from "comlink"; +import { z } from "zod"; import downloadManager from "../download"; import { cmpNewLib2, extractRawExif } from "../exif"; import { getAllLocalFiles, getLocalTrashedFiles } from "../files"; @@ -46,6 +47,46 @@ interface IndexableItem { remoteDerivedData: RemoteDerivedData | undefined; } +/** + * The port used to communicate with the Node.js ML worker process + * + * See: [Note: ML IPC] + * */ +let _port: MessagePort | undefined; + +globalThis.onmessage = (event: MessageEvent) => { + if (event.data == "createMLWorker/port") { + _port = event.ports[0]; + _port?.start(); + } +}; + +const IPCResponse = z.object({ + id: z.number(), + data: z.any(), +}); + +/** + * Our hand-rolled IPC handler and router - the web worker end. + * + * Sibling of the handleMessage function (in `ml-worker.ts`) in the desktop app. + */ +const electronMLWorker = async (type: string, data: string) => { + const port = ensure(_port); + // Generate a unique nonce to identify this RPC interaction. + const id = Math.random(); + return new Promise((resolve) => { + const handleMessage = (event: MessageEvent) => { + const response = IPCResponse.parse(event.data); + if (response.id != id) return; + port.removeEventListener("message", handleMessage); + resolve(response.data); + }; + port.addEventListener("message", handleMessage); + port.postMessage({ type, id, data }); + }); +}; + /** * Run operations related to machine learning (e.g. indexing) in a Web Worker. * @@ -113,6 +154,12 @@ export class MLWorker { // need to monkey patch it (This also ensures that it is not tree // shaken). globalThis.DOMParser = DOMParser; + + void (async () => { + console.log("yyy calling foo with 3"); + const res = await electronMLWorker("foo", "3"); + console.log("yyy calling foo with 3 result", res); + })(); } /** @@ -245,10 +292,6 @@ export class MLWorker { expose(MLWorker); -globalThis.onmessage = (event: MessageEvent) => { - console.log("worker onmessage", event); -}; - /** * Find out files which need to be indexed. Then index the next batch of them. *