Skip to content

Commit 2956ee4

Browse files
PauanPauan
committed
Adding in thread_pool::spawn function
1 parent 9396331 commit 2956ee4

File tree

3 files changed

+225
-184
lines changed

3 files changed

+225
-184
lines changed

‎sdk/src/polyfill/worker.ts

+144-172
Original file line numberDiff line numberDiff line change
@@ -1,216 +1,188 @@
1-
function patch($worker: typeof import("node:worker_threads"), $os: typeof import("node:os")) {
2-
// This is technically not a part of the Worker polyfill,
3-
// but Workers are used for multi-threading, so this is often
4-
// needed when writing Worker code.
5-
if (globalThis.navigator == null) {
6-
globalThis.navigator = {
7-
hardwareConcurrency: $os.cpus().length,
8-
} as Navigator;
9-
}
10-
11-
globalThis.Worker = class Worker extends EventTarget {
12-
private _worker: import("node:worker_threads").Worker;
1+
import * as $worker from "node:worker_threads";
2+
import * as $os from "node:os";
3+
4+
// This is technically not a part of the Worker polyfill,
5+
// but Workers are used for multi-threading, so this is often
6+
// needed when writing Worker code.
7+
if (globalThis.navigator == null) {
8+
globalThis.navigator = {
9+
hardwareConcurrency: $os.cpus().length,
10+
} as Navigator;
11+
}
1312

14-
constructor(url: string | URL, options?: WorkerOptions | undefined) {
15-
super();
13+
globalThis.Worker = class Worker extends EventTarget {
14+
private _worker: import("node:worker_threads").Worker;
1615

17-
if (url instanceof URL) {
18-
if (url.protocol !== "file:") {
19-
throw new Error("Worker only supports file: URLs");
20-
}
16+
constructor(url: string | URL, options?: WorkerOptions | undefined) {
17+
super();
2118

22-
url = url.href;
23-
24-
} else {
25-
throw new Error("Filepaths are unreliable, use `new URL(\"...\", import.meta.url)` instead.");
19+
if (url instanceof URL) {
20+
if (url.protocol !== "file:") {
21+
throw new Error("Worker only supports file: URLs");
2622
}
2723

28-
if (!options || options.type !== "module") {
29-
throw new Error("Workers must use \`type: \"module\"\`");
30-
}
24+
url = url.href;
3125

32-
// This uses some funky stuff like `patch.toString()`.
33-
//
34-
// This is needed so that it can synchronously run the polyfill code
35-
// inside of the worker.
36-
//
37-
// It can't use `require` because the file doesn't have a `.cjs` file extension.
38-
//
39-
// It can't use `import` because that's asynchronous, and the file path
40-
// might be different if using a bundler.
41-
const code = `
42-
${patch.toString()}
43-
44-
// Inject the polyfill into the worker
45-
patch(require("node:worker_threads"), require("node:os"));
46-
47-
const { workerData } = require("node:worker_threads");
48-
49-
// This actually loads and runs the worker file
50-
import(workerData.url)
51-
.catch((e) => {
52-
// TODO maybe it should send a message to the parent?
53-
console.error(e.stack);
54-
});
55-
`;
56-
57-
this._worker = new $worker.Worker(code, {
58-
eval: true,
59-
workerData: {
60-
url,
61-
},
62-
});
63-
64-
this._worker.on("message", (data) => {
65-
this.dispatchEvent(new MessageEvent("message", { data }));
66-
});
67-
68-
this._worker.on("messageerror", (error) => {
69-
throw new Error("UNIMPLEMENTED");
70-
});
71-
72-
this._worker.on("error", (error) => {
73-
// TODO attach the error to the event somehow
74-
const event = new Event("error");
75-
this.dispatchEvent(event);
76-
});
26+
} else {
27+
throw new Error("Filepaths are unreliable, use `new URL(\"...\", import.meta.url)` instead.");
7728
}
7829

79-
set onmessage(f: () => void) {
80-
throw new Error("UNIMPLEMENTED");
30+
if (!options || options.type !== "module") {
31+
throw new Error("Workers must use \`type: \"module\"\`");
8132
}
8233

83-
set onmessageerror(f: () => void) {
84-
throw new Error("UNIMPLEMENTED");
85-
}
34+
const code = `
35+
const { workerData } = require("node:worker_threads");
36+
37+
import(workerData.polyfill)
38+
.then(() => import(workerData.url))
39+
.catch((e) => {
40+
// TODO maybe it should send a message to the parent?
41+
console.error(e.stack);
42+
});
43+
`;
44+
45+
this._worker = new $worker.Worker(code, {
46+
eval: true,
47+
workerData: {
48+
url,
49+
polyfill: new URL("node-polyfill.js", import.meta.url).href,
50+
},
51+
});
8652

87-
set onerror(f: () => void) {
53+
this._worker.on("message", (data) => {
54+
this.dispatchEvent(new MessageEvent("message", { data }));
55+
});
56+
57+
this._worker.on("messageerror", (error) => {
8858
throw new Error("UNIMPLEMENTED");
89-
}
59+
});
9060

91-
postMessage(message: any, transfer: Array<Transferable>): void;
92-
postMessage(message: any, options?: StructuredSerializeOptions | undefined): void;
93-
postMessage(value: any, transfer: any) {
94-
this._worker.postMessage(value, transfer);
95-
}
61+
this._worker.on("error", (error) => {
62+
// TODO attach the error to the event somehow
63+
const event = new Event("error");
64+
this.dispatchEvent(event);
65+
});
66+
}
9667

97-
terminate() {
98-
this._worker.terminate();
99-
}
68+
set onmessage(f: () => void) {
69+
throw new Error("UNIMPLEMENTED");
70+
}
10071

101-
// This is Node-specific, it allows the process to exit
102-
// even if the Worker is still running.
103-
unref() {
104-
this._worker.unref();
105-
}
106-
};
72+
set onmessageerror(f: () => void) {
73+
throw new Error("UNIMPLEMENTED");
74+
}
10775

76+
set onerror(f: () => void) {
77+
throw new Error("UNIMPLEMENTED");
78+
}
10879

109-
if (!$worker.isMainThread) {
110-
const globals = globalThis as unknown as DedicatedWorkerGlobalScope;
80+
postMessage(message: any, transfer: Array<Transferable>): void;
81+
postMessage(message: any, options?: StructuredSerializeOptions | undefined): void;
82+
postMessage(value: any, transfer: any) {
83+
this._worker.postMessage(value, transfer);
84+
}
11185

112-
// This is used to create the onmessage, onmessageerror, and onerror setters
113-
const makeSetter = (prop: string, event: string) => {
114-
let oldvalue: () => void;
86+
terminate() {
87+
this._worker.terminate();
88+
}
11589

116-
Object.defineProperty(globals, prop, {
117-
get() {
118-
return oldvalue;
119-
},
120-
set(value) {
121-
if (oldvalue) {
122-
globals.removeEventListener(event, oldvalue);
123-
}
90+
// This is Node-specific, it allows the process to exit
91+
// even if the Worker is still running.
92+
unref() {
93+
this._worker.unref();
94+
}
95+
};
12496

125-
oldvalue = value;
12697

127-
if (oldvalue) {
128-
globals.addEventListener(event, oldvalue);
129-
}
130-
},
131-
});
132-
};
98+
if (!$worker.isMainThread) {
99+
const globals = globalThis as unknown as DedicatedWorkerGlobalScope;
133100

134-
// This makes sure that `f` is only run once
135-
const memoize = (f: () => void) => {
136-
let run = false;
101+
// This is used to create the onmessage, onmessageerror, and onerror setters
102+
const makeSetter = (prop: string, event: string) => {
103+
let oldvalue: () => void;
137104

138-
return () => {
139-
if (!run) {
140-
run = true;
141-
f();
105+
Object.defineProperty(globals, prop, {
106+
get() {
107+
return oldvalue;
108+
},
109+
set(value) {
110+
if (oldvalue) {
111+
globals.removeEventListener(event, oldvalue);
142112
}
143-
};
144-
};
145113

114+
oldvalue = value;
146115

147-
// We only start listening for messages / errors when the worker calls addEventListener
148-
const startOnMessage = memoize(() => {
149-
$worker.parentPort!.on("message", (data) => {
150-
workerEvents.dispatchEvent(new MessageEvent("message", { data }));
151-
});
116+
if (oldvalue) {
117+
globals.addEventListener(event, oldvalue);
118+
}
119+
},
152120
});
121+
};
153122

154-
const startOnMessageError = memoize(() => {
155-
throw new Error("UNIMPLEMENTED");
156-
});
123+
// This makes sure that `f` is only run once
124+
const memoize = (f: () => void) => {
125+
let run = false;
157126

158-
const startOnError = memoize(() => {
159-
$worker.parentPort!.on("error", (data) => {
160-
workerEvents.dispatchEvent(new Event("error"));
161-
});
127+
return () => {
128+
if (!run) {
129+
run = true;
130+
f();
131+
}
132+
};
133+
};
134+
135+
136+
// We only start listening for messages / errors when the worker calls addEventListener
137+
const startOnMessage = memoize(() => {
138+
$worker.parentPort!.on("message", (data) => {
139+
workerEvents.dispatchEvent(new MessageEvent("message", { data }));
162140
});
141+
});
163142

143+
const startOnMessageError = memoize(() => {
144+
throw new Error("UNIMPLEMENTED");
145+
});
164146

165-
// Node workers don't have top-level events, so we have to make our own
166-
const workerEvents = new EventTarget();
147+
const startOnError = memoize(() => {
148+
$worker.parentPort!.on("error", (data) => {
149+
workerEvents.dispatchEvent(new Event("error"));
150+
});
151+
});
167152

168-
globals.close = () => {
169-
process.exit();
170-
};
171153

172-
globals.addEventListener = (type: string, callback: EventListenerOrEventListenerObject | null, options?: boolean | EventListenerOptions | undefined) => {
173-
workerEvents.addEventListener(type, callback, options);
154+
// Node workers don't have top-level events, so we have to make our own
155+
const workerEvents = new EventTarget();
174156

175-
if (type === "message") {
176-
startOnMessage();
177-
} else if (type === "messageerror") {
178-
startOnMessageError();
179-
} else if (type === "error") {
180-
startOnError();
181-
}
182-
};
157+
globals.close = () => {
158+
process.exit();
159+
};
183160

184-
globals.removeEventListener = (type: string, callback: EventListenerOrEventListenerObject | null, options?: boolean | EventListenerOptions | undefined) => {
185-
workerEvents.removeEventListener(type, callback, options);
186-
};
161+
globals.addEventListener = (type: string, callback: EventListenerOrEventListenerObject | null, options?: boolean | EventListenerOptions | undefined) => {
162+
workerEvents.addEventListener(type, callback, options);
187163

188-
function postMessage(message: any, transfer: Transferable[]): void;
189-
function postMessage(message: any, options?: StructuredSerializeOptions | undefined): void;
190-
function postMessage(value: any, transfer: any) {
191-
$worker.parentPort!.postMessage(value, transfer);
164+
if (type === "message") {
165+
startOnMessage();
166+
} else if (type === "messageerror") {
167+
startOnMessageError();
168+
} else if (type === "error") {
169+
startOnError();
192170
}
171+
};
193172

194-
globals.postMessage = postMessage;
173+
globals.removeEventListener = (type: string, callback: EventListenerOrEventListenerObject | null, options?: boolean | EventListenerOptions | undefined) => {
174+
workerEvents.removeEventListener(type, callback, options);
175+
};
195176

196-
makeSetter("onmessage", "message");
197-
makeSetter("onmessageerror", "messageerror");
198-
makeSetter("onerror", "error");
177+
function postMessage(message: any, transfer: Transferable[]): void;
178+
function postMessage(message: any, options?: StructuredSerializeOptions | undefined): void;
179+
function postMessage(value: any, transfer: any) {
180+
$worker.parentPort!.postMessage(value, transfer);
199181
}
200-
}
201-
202182

203-
async function polyfill() {
204-
const [$worker, $os] = await Promise.all([
205-
import("node:worker_threads"),
206-
import("node:os"),
207-
]);
183+
globals.postMessage = postMessage;
208184

209-
patch($worker, $os);
185+
makeSetter("onmessage", "message");
186+
makeSetter("onmessageerror", "messageerror");
187+
makeSetter("onerror", "error");
210188
}
211-
212-
if (globalThis.Worker == null) {
213-
await polyfill();
214-
}
215-
216-
export {};

‎wasm/src/lib.rs

+14-4
Original file line numberDiff line numberDiff line change
@@ -167,10 +167,20 @@ pub use types::Field;
167167
#[cfg(not(test))]
168168
mod thread_pool;
169169

170-
use wasm_bindgen::prelude::*;
170+
#[cfg(test)]
171+
mod thread_pool {
172+
use std::future::Future;
173+
174+
pub fn spawn<A, F>(f: F) -> impl Future<Output = A>
175+
where
176+
A: Send + 'static,
177+
F: FnOnce() -> A + Send + 'static,
178+
{
179+
async move { f() }
180+
}
181+
}
171182

172-
#[cfg(not(test))]
173-
use thread_pool::ThreadPool;
183+
use wasm_bindgen::prelude::*;
174184

175185
use std::str::FromStr;
176186

@@ -219,7 +229,7 @@ use types::native;
219229
pub async fn init_thread_pool(url: web_sys::Url, num_threads: usize) -> Result<(), JsValue> {
220230
console_error_panic_hook::set_once();
221231

222-
ThreadPool::builder().url(url).num_threads(num_threads).build_global().await?;
232+
thread_pool::ThreadPool::builder().url(url).num_threads(num_threads).build_global().await?;
223233

224234
Ok(())
225235
}

0 commit comments

Comments
 (0)