Desktop NodeJS worker and NODEFS support

.github/workflows/main.yml

@@ -257,6 +257,8 @@ jobs:
         cache: 'yarn'
     - run: yarn config set network-timeout 600000  # `yarn install` often takes so long time on in the Windows env.
     - run: yarn install --frozen-lockfile --ignore-scripts  # `--ignore-scripts` is necessary because the `postinstall` script of the desktop app sample project fails due to https://github.com/yarnpkg/yarn/issues/7694
+    - run: make common
+      working-directory: .
     - name: Lint
       if: ${{ matrix.os == 'ubuntu-latest' }}  # The glob pattern passed to ESLint is hardcoded as POSIX, so it does not work on Windows.
       run: |

packages/desktop/bin-src/dump_artifacts/index.ts

@@ -293,8 +293,10 @@ async function dumpManifest(options: DumpManifestOptions) {
 
   // TODO: Runtime type validation
   const manifestData: DesktopAppManifest = {
-    embed: stliteManifest.embed || false,
+    embed: stliteManifest.embed ?? false,
     idbfsMountpoints: stliteManifest.idbfsMountpoints,
+    nodeJsWorker: stliteManifest.nodeJsWorker ?? false,
+    nodefsMountpoints: stliteManifest.nodefsMountpoints,
   };
 
   const manifestDataStr = JSON.stringify(manifestData, null, 2);

packages/desktop/electron/main.ts

@@ -1,6 +1,7 @@
 import { app, BrowserWindow, ipcMain, protocol } from "electron";
 import * as path from "path";
 import * as fsPromises from "fs/promises";
+import workerThreads from "node:worker_threads";
 import { walkRead } from "./file";
 
 if (process.env.NODE_ENV === "development") {
@@ -18,6 +19,8 @@ if (process.env.NODE_ENV === "development") {
 export interface DesktopAppManifest {
   embed: boolean;
   idbfsMountpoints: string[] | undefined;
+  nodeJsWorker: boolean;
+  nodefsMountpoints: Record<string, string> | undefined;
 }
 async function readManifest(): Promise<DesktopAppManifest> {
   const manifestPath = path.resolve(__dirname, "../stlite-manifest.json");
@@ -30,21 +33,31 @@ async function readManifest(): Promise<DesktopAppManifest> {
   return {
     embed: maybeManifestData.embed ?? false,
     idbfsMountpoints: maybeManifestData.idbfsMountpoints,
+    nodeJsWorker: maybeManifestData.nodeJsWorker ?? false,
+    nodefsMountpoints: maybeManifestData.nodefsMountpoints,
   };
 }
 
 const createWindow = async () => {
   const manifest = await readManifest();
 
+  const additionalArguments: string[] = [];
+  if (manifest.idbfsMountpoints) {
+    additionalArguments.push(
+      `--idbfs-mountpoints=${JSON.stringify(manifest.idbfsMountpoints)}`
+    );
+  }
+  if (manifest.nodeJsWorker) {
+    additionalArguments.push("--nodejs-worker");
+  }
+
   const mainWindow = new BrowserWindow({
     width: 1280,
     height: 720,
     webPreferences: {
       preload: path.join(__dirname, "preload.js"),
       sandbox: true, // https://www.electronjs.org/docs/latest/tutorial/security#4-enable-process-sandboxing
-      additionalArguments: manifest.idbfsMountpoints
-        ? [`--idbfs-mountpoints=${JSON.stringify(manifest.idbfsMountpoints)}`]
-        : undefined,
+      additionalArguments,
     },
   });
 
@@ -113,6 +126,61 @@ const createWindow = async () => {
     ipcMain.removeHandler("readStreamlitAppDirectory");
   });
 
+  let worker: workerThreads.Worker | null = null;
+  ipcMain.handle("initializeNodeJsWorker", (ev) => {
+    if (!isValidIpcSender(ev.senderFrame)) {
+      throw new Error("Invalid IPC sender");
+    }
+
+    // Use the ESM version of Pyodide because `importScripts()` can't be used in this environment.
+    const defaultPyodideUrl = path.resolve(__dirname, "../pyodide/pyodide.mjs");
+
+    function onMessageFromWorker(value: any) {
+      mainWindow.webContents.send("messageFromNodeJsWorker", value);
+    }
+    worker = new workerThreads.Worker(path.resolve(__dirname, "./worker.js"), {
+      env: {
+        PYODIDE_URL: defaultPyodideUrl,
+        NODEFS_MOUNTPOINTS: JSON.stringify(manifest.nodefsMountpoints),
+      },
+    });
+    worker.on("message", (value) => {
+      onMessageFromWorker(value);
+    });
+  });
+  ipcMain.on("messageToNodeJsWorker", (ev, { data, portId }) => {
+    if (!isValidIpcSender(ev.senderFrame)) {
+      throw new Error("Invalid IPC sender");
+    }
+
+    if (worker == null) {
+      return;
+    }
+
+    const channel = new workerThreads.MessageChannel();
+
+    channel.port1.on("message", (e) => {
+      ev.reply(`nodeJsWorker-portMessage-${portId}`, e);
+    });
+
+    const eventSim = { data, port: channel.port2 };
+    worker.postMessage(eventSim, [channel.port2]);
+  });
+  ipcMain.handle("terminate", (ev, { data, portId }) => {
+    if (!isValidIpcSender(ev.senderFrame)) {
+      throw new Error("Invalid IPC sender");
+    }
+
+    worker?.terminate();
+    worker = null;
+  });
+
+  mainWindow.on("closed", () => {
+    ipcMain.removeHandler("initializeNodeJsWorker");
+    ipcMain.removeHandler("messageToNodeJsWorker");
+    ipcMain.removeHandler("terminate");
+  });
+
   // Even when the entrypoint is a local file like the production build,
   // we use .loadURL() with an absolute URL with the `file://` schema
   // instead of passing a file path to .loadFile()

packages/desktop/electron/preload.ts

@@ -17,3 +17,39 @@ contextBridge.exposeInMainWorld("archives", {
   readStreamlitAppDirectory: () =>
     ipcRenderer.invoke("readStreamlitAppDirectory"),
 });
+
+function getRandomInt() {
+  return Math.floor(Math.random() * 1000000);
+}
+
+const nodeJsWorkerAPI = {
+  USE_NODEJS_WORKER: process.argv.includes("--nodejs-worker"),
+  initialize: () => ipcRenderer.invoke("initializeNodeJsWorker"),
+  postMessage: ({
+    data,
+    onPortMessage,
+  }: {
+    data: any;
+    onPortMessage: ((arg: any) => void) | null;
+  }) => {
+    console.debug("nodeJsWorkerAPI.postMessage", { data, onPortMessage });
+    // When the `contextIsolation` is enabled, `MessagePort` objects cannot be transferred between contexts even with ipcRenderer.postMessage(),
+    // so we need to simulate the `MessagePort` API with `ipcRenderer.on` and `ipcRenderer.send`.
+    // This works in a combination with the `NodeJsWorkerMock` class in `nodejs-worker.ts`.
+    const portId = onPortMessage && getRandomInt();
+    ipcRenderer.send("messageToNodeJsWorker", { data, portId });
+    if (portId) {
+      ipcRenderer.on(`nodeJsWorker-portMessage-${portId}`, (_event, arg) => {
+        onPortMessage(arg);
+      });
+    }
+  },
+  onMessage: (callback) =>
+    ipcRenderer.on("messageFromNodeJsWorker", (_event, value) => {
+      console.debug("nodeJsWorkerAPI.onMessage", value);
+      callback(value);
+    }),
+  terminate: () => ipcRenderer.invoke("terminateNodeJsWorker"),
+};
+contextBridge.exposeInMainWorld("nodeJsWorkerAPI", nodeJsWorkerAPI);
+export type NodeJsWorkerAPI = typeof nodeJsWorkerAPI;

packages/desktop/electron/worker.ts

@@ -0,0 +1,25 @@
+import { parentPort } from "node:worker_threads";
+import { startWorkerEnv } from "@stlite/kernel/src/worker-runtime";
+
+function postMessage(value: any) {
+  console.debug("[worker thread] postMessage from worker", value);
+  parentPort?.postMessage(value);
+}
+
+// TODO: Runtime type validation
+const nodefsMountpoints =
+  process.env.NODEFS_MOUNTPOINTS && JSON.parse(process.env.NODEFS_MOUNTPOINTS);
+
+const handleMessage = startWorkerEnv(
+  process.env.PYODIDE_URL as string,
+  postMessage,
+  {
+    nodefsMountpoints,
+  }
+);
+
+parentPort?.on("message", ({ data, port }) => {
+  console.debug("[worker thread] parentPort.onMessage", { data, port });
+  const simEvent = { data, ports: [port] };
+  handleMessage(simEvent as any);
+});

packages/desktop/package.json

@@ -102,9 +102,10 @@
   "stlite": {
     "desktop": {
       "embed": true,
-      "idbfsMountpoints": [
-        "/mnt"
-      ]
+      "nodeJsWorker": true,
+      "nodefsMountpoints": {
+        "/mnt": "."
+      }
     }
   }
 }

packages/desktop/pyodide-files.txt

@@ -1,5 +1,5 @@
 pyodide/pyodide.asm.js
 pyodide/pyodide.asm.wasm
-pyodide/pyodide.js
+pyodide/pyodide.mjs
 pyodide/python_stdlib.zip
 pyodide/pyodide-lock.json

packages/desktop/scripts/build_electron.js

@@ -12,6 +12,7 @@ const production = process.env.NODE_ENV === "production";
   entryPoints: [
     path.resolve(__dirname, "../electron/main.ts"),
     path.resolve(__dirname, "../electron/preload.ts"),
+    path.resolve(__dirname, "../electron/worker.ts"),
   ],
   bundle: true,
   minify: production,

packages/desktop/src/App.tsx

@@ -2,18 +2,19 @@ import React, { useState, useEffect } from "react";
 import { StliteKernel, StliteKernelOptions } from "@stlite/kernel";
 import StreamlitApp from "./StreamlitApp";
 import { makeToastKernelCallbacks } from "@stlite/common-react";
+import { USE_NODEJS_WORKER, NodeJsWorkerMock } from "./nodejs-worker";
 import "@stlite/common-react/src/toastify-components/toastify.css";
 
 let pyodideUrl: string | undefined;
-if (process.env.NODE_ENV === "production") {
+if (process.env.NODE_ENV === "production" && !USE_NODEJS_WORKER) {
   // The `pyodide` directory including `pyodide.js` is downloaded
   // to the build target directory at the build time for production release.
   // See the "build:pyodide" NPM script.
   // Ref: https://pyodide.org/en/stable/usage/downloading-and-deploying.html
   // We set the path here to be loaded in the worker via `importScript()`.
   const currentURL = window.location.href;
   const parentURL = currentURL.split("/").slice(0, -1).join("/") + "/";
-  pyodideUrl = parentURL + "pyodide/pyodide.js";
+  pyodideUrl = parentURL + "pyodide/pyodide.mjs";
 }
 
 function App() {
@@ -55,6 +56,9 @@ function App() {
           mountedSitePackagesSnapshotFilePath,
           pyodideUrl,
           idbfsMountpoints: window.appConfig.idbfsMountpoints,
+          worker: USE_NODEJS_WORKER
+            ? (new NodeJsWorkerMock() as unknown as Worker)
+            : undefined,
           ...makeToastKernelCallbacks(),
         });
         setKernel(kernel);

packages/desktop/src/global.d.ts

@@ -1,4 +1,4 @@
-import type { AppConfig } from "../electron/preload";
+import type { AppConfig, NodeJsWorkerAPI } from "../electron/preload";
 
 export declare global {
   interface Window {
@@ -8,5 +8,6 @@ export declare global {
       readRequirements: () => Promise<string[]>;
       readStreamlitAppDirectory: () => Promise<Record<string, Buffer>>;
     };
+    nodeJsWorkerAPI: NodeJsWorkerAPI;
   }
 }

packages/desktop/src/nodejs-worker.ts

@@ -0,0 +1,36 @@
+export const USE_NODEJS_WORKER = window.nodeJsWorkerAPI.USE_NODEJS_WORKER;
+
+interface MessageEventLike<T = any> {
+  readonly data: T;
+}
+
+export class NodeJsWorkerMock {
+  private initializePromise: Promise<void>;
+
+  constructor() {
+    this.initializePromise = window.nodeJsWorkerAPI.initialize();
+
+    window.nodeJsWorkerAPI.onMessage((data) => {
+      this.onmessage && this.onmessage({ data });
+    });
+  }
+
+  postMessage(data: any, transfer?: [MessagePort]) {
+    this.initializePromise.then(() => {
+      const port = transfer ? transfer[0] : null;
+      const onPortMessage =
+        port &&
+        ((arg: any) => {
+          port.postMessage(arg);
+        });
+
+      window.nodeJsWorkerAPI.postMessage({ data, onPortMessage });
+    });
+  }
+
+  onmessage: ((e: MessageEventLike) => void) | null = null;
+
+  terminate() {
+    window.nodeJsWorkerAPI.terminate();
+  }
+}

packages/kernel/src/declarations.d.ts

@@ -8,10 +8,3 @@ declare module "*.whl" {
   const res: string;
   return res;
 }
-
-// Declarations for the `self` object in worker.ts some of whose properties are used to share values with the Python environment.
-interface Window {
-  __logCallback__: (levelno: number, msg: string) => void;
-  __streamlitFlagOptions__: Record<string, PyodideConvertiblePrimitive>;
-  __scriptFinishedCallback__: () => void;
-}

packages/kernel/src/kernel.ts

@@ -115,6 +115,12 @@ export interface StliteKernelOptions {
   onLoad?: () => void;
 
   onError?: (error: Error) => void;
+
+  /**
+   * The worker to be used, which can be optionally passed.
+   * Desktop apps with NodeJS-backed worker is one of the use cases.
+   */
+  worker?: globalThis.Worker;
 }
 
 export class StliteKernel {
@@ -145,10 +151,14 @@ export class StliteKernel {
     this.onLoad = options.onLoad;
     this.onError = options.onError;
 
-    // HACK: Use `CrossOriginWorkerMaker` imported as `Worker` here.
-    // Read the comment in `cross-origin-worker.ts` for the detail.
-    const workerMaker = new Worker(new URL("./worker.js", import.meta.url));
-    this._worker = workerMaker.worker;
+    if (options.worker) {
+      this._worker = options.worker;
+    } else {
+      // HACK: Use `CrossOriginWorkerMaker` imported as `Worker` here.
+      // Read the comment in `cross-origin-worker.ts` for the detail.
+      const workerMaker = new Worker(new URL("./worker.js", import.meta.url));
+      this._worker = workerMaker.worker;
+    }
 
     this._worker.onmessage = (e) => {
       this._processWorkerMessage(e.data);

packages/kernel/src/types.ts

@@ -52,6 +52,7 @@
   mountedSitePackagesSnapshotFilePath?: string;
   streamlitConfig?: StreamlitConfig;
   idbfsMountpoints?: string[];
+  nodefsMountpoints?: Record<string, string>;
 }
 
 /**
@@ -88,7 +89,7 @@
   data: {
     path: string;
     data: string | ArrayBufferView;
     opts?: Record<string, any>;
   };
 }
 export interface InMessageFileRename extends InMessageBase {
@@ -169,7 +170,7 @@
 interface ReplyMessageBase {
   type: string;
   error?: Error;
   data?: any;
 }
 export interface ReplyMessageHttpResponse extends ReplyMessageBase {
   type: "http:response";