diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 05642fc95f..cc1ae5764c 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -11,8 +11,11 @@ import { app, utilityProcess } from "electron/main"; import path from "node:path"; import log from "../log"; +/** The active ML worker (utility) process, if any. */ +let _child: UtilityProcess | undefined; + /** - * Create a new ML worker process. + * Create a new ML worker process, terminating the older ones (if any). * * [Note: ML IPC] * @@ -68,6 +71,12 @@ import log from "../log"; * to be relayed using `postMessage`. */ export const createMLWorker = (window: BrowserWindow) => { + if (_child) { + log.debug(() => "Terminating previous ML worker process"); + _child.kill(); + _child = undefined; + } + const { port1, port2 } = new MessageChannelMain(); const child = utilityProcess.fork(path.join(__dirname, "ml-worker.js")); @@ -77,6 +86,8 @@ export const createMLWorker = (window: BrowserWindow) => { window.webContents.postMessage("createMLWorker/port", undefined, [port2]); handleMessagesFromUtilityProcess(child); + + _child = child; }; /** diff --git a/web/packages/base/types/ipc.ts b/web/packages/base/types/ipc.ts index 748490ed58..c0644760c0 100644 --- a/web/packages/base/types/ipc.ts +++ b/web/packages/base/types/ipc.ts @@ -335,7 +335,7 @@ export interface Electron { // - ML /** - * Create a new ML worker. + * Create a new ML worker, terminating the older ones (if any). * * This creates a new Node.js utility process, and sets things up so that we * can communicate directly with that utility process using a diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index ba8083ad74..5b57dade21 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -98,22 +98,27 @@ export const terminateMLWorker = async () => { * Obtain a port from the Node.js layer that can be used to communicate with the * ML worker process. */ -const createMLWorker = async (electron: Electron): Promise => { - electron.createMLWorker(); - +const createMLWorker = (electron: Electron): Promise => { // The main process will do its thing, and send back the port it created to // us by sending an message on the "createMLWorker/port" channel via the // postMessage API. This roundabout way is needed because MessagePorts // cannot be transferred via the usual send/invoke pattern. - return new Promise((resolve) => { - window.onmessage = ({ source, data, ports }: MessageEvent) => { - // The source check verifies that the message is coming from the + const port = new Promise((resolve) => { + const l = ({ source, data, ports }: MessageEvent) => { + // The source check verifies that the message is coming from our own // preload script. The data is the message that was posted. - if (source == window && data == "createMLWorker/port") + if (source == window && data == "createMLWorker/port") { + window.removeEventListener("message", l); resolve(ensure(ports[0])); + } }; + window.addEventListener("message", l); }); + + electron.createMLWorker(); + + return port; }; /**