@@ -3,26 +3,30 @@ import { WASI as NodeWASI } from "wasi" |
3 | 3 | import { WASI as MicroWASI, useAll } from "uwasi" |
4 | 4 | import * as fs from "fs/promises" |
5 | 5 | import path from "path"; |
| 6 | +import { Worker, parentPort } from "node:worker_threads"; |
6 | 7 |
|
7 | 8 | const WASI = { |
8 | | - MicroWASI: ({ programName }) => { |
| 9 | + MicroWASI: ({ args }) => { |
9 | 10 | const wasi = new MicroWASI({ |
10 | | - args: [path.basename(programName)], |
| 11 | + args: args, |
11 | 12 | env: {}, |
12 | 13 | features: [useAll()], |
13 | 14 | }) |
14 | 15 |
|
15 | 16 | return { |
16 | 17 | wasiImport: wasi.wasiImport, |
| 18 | + setInstance(instance) { |
| 19 | + wasi.instance = instance; |
| 20 | + }, |
17 | 21 | start(instance, swift) { |
18 | 22 | wasi.initialize(instance); |
19 | 23 | swift.main(); |
20 | 24 | } |
21 | 25 | } |
22 | 26 | }, |
23 | | - Node: ({ programName }) => { |
| 27 | + Node: ({ args }) => { |
24 | 28 | const wasi = new NodeWASI({ |
25 | | - args: [path.basename(programName)], |
| 29 | + args: args, |
26 | 30 | env: {}, |
27 | 31 | preopens: { |
28 | 32 | "/": "./", |
@@ -44,12 +48,9 @@ const WASI = { |
44 | 48 | const selectWASIBackend = () => { |
45 | 49 | const value = process.env["JAVASCRIPTKIT_WASI_BACKEND"] |
46 | 50 | if (value) { |
47 | | - const backend = WASI[value]; |
48 | | - if (backend) { |
49 | | - return backend; |
50 | | - } |
| 51 | + return value; |
51 | 52 | } |
52 | | - return WASI.Node; |
| 53 | + return "Node" |
53 | 54 | }; |
54 | 55 |
|
55 | 56 | function isUsingSharedMemory(module) { |
@@ -62,33 +63,125 @@ function isUsingSharedMemory(module) { |
62 | 63 | return false; |
63 | 64 | } |
64 | 65 |
|
65 | | -export const startWasiTask = async (wasmPath, wasiConstructor = selectWASIBackend()) => { |
66 | | - const swift = new SwiftRuntime(); |
67 | | - // Fetch our Wasm File |
68 | | - const wasmBinary = await fs.readFile(wasmPath); |
69 | | - const wasi = wasiConstructor({ programName: wasmPath }); |
70 | | - |
71 | | - const module = await WebAssembly.compile(wasmBinary); |
72 | | - |
73 | | - const importObject = { |
| 66 | +function constructBaseImportObject(wasi, swift) { |
| 67 | + return { |
74 | 68 | wasi_snapshot_preview1: wasi.wasiImport, |
75 | | - javascript_kit: swift.importObjects(), |
| 69 | + javascript_kit: swift.wasmImports, |
76 | 70 | benchmark_helper: { |
77 | 71 | noop: () => {}, |
78 | 72 | noop_with_int: (_) => {}, |
| 73 | + }, |
| 74 | + } |
| 75 | +} |
| 76 | + |
| 77 | +export async function startWasiChildThread(event) { |
| 78 | + const { module, programName, memory, tid, startArg } = event; |
| 79 | + const swift = new SwiftRuntime({ |
| 80 | + sharedMemory: true, |
| 81 | + threadChannel: { |
| 82 | + wakeUpMainThread: parentPort.postMessage.bind(parentPort), |
| 83 | + listenWakeEventFromMainThread: (listener) => { |
| 84 | + parentPort.on("message", listener) |
| 85 | + } |
| 86 | + } |
| 87 | + }); |
| 88 | + // Use uwasi for child threads because Node.js WASI cannot be used without calling |
| 89 | + // `WASI.start` or `WASI.initialize`, which is already called in the main thread and |
| 90 | + // will cause an error if called again. |
| 91 | + const wasi = WASI.MicroWASI({ programName }); |
| 92 | + |
| 93 | + const importObject = constructBaseImportObject(wasi, swift); |
| 94 | + |
| 95 | + importObject["wasi"] = { |
| 96 | + "thread-spawn": () => { |
| 97 | + throw new Error("Cannot spawn a new thread from a worker thread") |
79 | 98 | } |
80 | 99 | }; |
| 100 | + importObject["env"] = { memory }; |
| 101 | + importObject["JavaScriptEventLoopTestSupportTests"] = { |
| 102 | + "isMainThread": () => false, |
| 103 | + } |
| 104 | + |
| 105 | + const instance = await WebAssembly.instantiate(module, importObject); |
| 106 | + swift.setInstance(instance); |
| 107 | + wasi.setInstance(instance); |
| 108 | + swift.startThread(tid, startArg); |
| 109 | +} |
| 110 | + |
| 111 | +class ThreadRegistry { |
| 112 | + workers = new Map(); |
| 113 | + nextTid = 1; |
| 114 | + |
| 115 | + spawnThread(module, programName, memory, startArg) { |
| 116 | + const tid = this.nextTid++; |
| 117 | + const selfFilePath = new URL(import.meta.url).pathname; |
| 118 | + const worker = new Worker(` |
| 119 | + const { parentPort } = require('node:worker_threads'); |
| 120 | +
|
| 121 | + Error.stackTraceLimit = 100; |
| 122 | + parentPort.once("message", async (event) => { |
| 123 | + const { selfFilePath } = event; |
| 124 | + const { startWasiChildThread } = await import(selfFilePath); |
| 125 | + await startWasiChildThread(event); |
| 126 | + }) |
| 127 | + `, { type: "module", eval: true }) |
| 128 | + |
| 129 | + worker.on("error", (error) => { |
| 130 | + console.error(`Worker thread ${tid} error:`, error); |
| 131 | + }); |
| 132 | + this.workers.set(tid, worker); |
| 133 | + worker.postMessage({ selfFilePath, module, programName, memory, tid, startArg }); |
| 134 | + return tid; |
| 135 | + } |
| 136 | + |
| 137 | + worker(tid) { |
| 138 | + return this.workers.get(tid); |
| 139 | + } |
| 140 | + |
| 141 | + wakeUpWorkerThread(tid) { |
| 142 | + const worker = this.workers.get(tid); |
| 143 | + worker.postMessage(null); |
| 144 | + } |
| 145 | +} |
| 146 | + |
| 147 | +export const startWasiTask = async (wasmPath, wasiConstructorKey = selectWASIBackend()) => { |
| 148 | + // Fetch our Wasm File |
| 149 | + const wasmBinary = await fs.readFile(wasmPath); |
| 150 | + const programName = wasmPath; |
| 151 | + const args = [path.basename(programName)]; |
| 152 | + args.push(...process.argv.slice(3)); |
| 153 | + const wasi = WASI[wasiConstructorKey]({ args }); |
| 154 | + |
| 155 | + const module = await WebAssembly.compile(wasmBinary); |
| 156 | + |
| 157 | + const sharedMemory = isUsingSharedMemory(module); |
| 158 | + const threadRegistry = new ThreadRegistry(); |
| 159 | + const swift = new SwiftRuntime({ |
| 160 | + sharedMemory, |
| 161 | + threadChannel: { |
| 162 | + wakeUpWorkerThread: threadRegistry.wakeUpWorkerThread.bind(threadRegistry), |
| 163 | + listenMainJobFromWorkerThread: (tid, listener) => { |
| 164 | + const worker = threadRegistry.worker(tid); |
| 165 | + worker.on("message", listener); |
| 166 | + } |
| 167 | + } |
| 168 | + }); |
| 169 | + |
| 170 | + const importObject = constructBaseImportObject(wasi, swift); |
| 171 | + |
| 172 | + importObject["JavaScriptEventLoopTestSupportTests"] = { |
| 173 | + "isMainThread": () => true, |
| 174 | + } |
81 | 175 |
|
82 | | - if (isUsingSharedMemory(module)) { |
83 | | - importObject["env"] = { |
84 | | - // We don't have JS API to get memory descriptor of imported memory |
85 | | - // at this moment, so we assume 256 pages (16MB) memory is enough |
86 | | - // large for initial memory size. |
87 | | - memory: new WebAssembly.Memory({ initial: 256, maximum: 16384, shared: true }), |
88 | | - }; |
| 176 | + if (sharedMemory) { |
| 177 | + // We don't have JS API to get memory descriptor of imported memory |
| 178 | + // at this moment, so we assume 256 pages (16MB) memory is enough |
| 179 | + // large for initial memory size. |
| 180 | + const memory = new WebAssembly.Memory({ initial: 256, maximum: 16384, shared: true }) |
| 181 | + importObject["env"] = { memory }; |
89 | 182 | importObject["wasi"] = { |
90 | | - "thread-spawn": () => { |
91 | | - throw new Error("thread-spawn not implemented"); |
| 183 | + "thread-spawn": (startArg) => { |
| 184 | + return threadRegistry.spawnThread(module, programName, memory, startArg); |
92 | 185 | } |
93 | 186 | } |
94 | 187 | } |
|