WaitGroups in JavaScript
Signalling to the main thread without message passing
So far in my JavaScript threading series, we’ve covered spawning threads & passing messages, mutexes, and condition variables. We can now start to have signaling and synchronization between threads, and we could start trying to make higher-level abstractions on-top of what we have, such as channels (which is something I did in C++ a while back).
However, before we do, there is one additional primitive I want to introduce, and that is Go-style WaitGroups. Currently, we are still using message passing in our examples to tell the main thread that our thread’s work is “done” and now they can read the value. Which really undermines the whole reason we went through all the trouble to share memory in the first place. If we’re still going to pass messages to signal the main thread1, then have we gained anything from sharing memory?
We could use a Condition Variable to signal things - but then we would also need to manage a mutex as well. Instead, what we want is a signaling primitive that is stand-alone - no mutex needed.
This is where Go’s WaitGroups come in. They are basically a waitable-counter. The main thread sets the initial value of the counter, and then waits until that counter hits zero. All of the other threads decrement the counter once they’re done with their work.
Here’s an example of how this would look in JavaScript:
// This would be in your main code
async function mainThread() {
// Setup our memory and wait group
const sab = new SharedArrayBuffer(Int32Array.BYTES_PER_ELEMENT * 3)
const ints = new Int32Array(sab)
const wg = new WaitGroup(ints, 0)
// Initialize our thread
const worker = new Worker('my-worker.js')
// Setup our work
wg.add(1)
worker.postMessage({__type: 'square', wg: 0, mem: ints, input: 3, dest: 1})
// Setup another work item
wg.add(1)
worker.postMessage({__type: 'cube', wg: 0, mem: ints, input: 4, dest: 2})
// Wait for the work to get done
await wg.waitAsync()
// Read our results
console.log('Square of 3: ', ints.at(1))
console.log('Cube of 4: ', ints.at(2))
}// my-worker.js
onmessage = (msg) => {
if (msg.__type === 'square') {
const v = msg.input
msg.mem.set([v * v], msg.dest)
new WaitGroup(msg.mem, msg.wg).done()
}
else if (msg.__type === 'cube') {
const v = msg.input
msg.mem.set([v * v * v], msg.dest)
new WaitGroup(msg.mem, msg.wg).done()
}
}With this code, we can now send tasks off to the worker thread, tell it where we want the results written, and then wait for the tasks to get done. Once the tasks are done, we simply read the memory. Everything is also self-contained, in the message2, so we can simply
In a way, this is simpler than the outstanding work map that stored Promise resolvers which we had before. All of the control flow for the main thread is linear. There is no magic data structure somewhere to correlate sent messages with received messages, and there’s no weird onmessage handler which then uses said data structure. Instead, we send some work, we wait for it, and we read the result. Very beautiful.
On the worker-side, it’s also fairly simple. We just call done when we’re done, rather than post a response back. We don’t have to parse and propagate a message id either. We call done and it takes care of the rest.
Actually Making a WaitGroup
So, now that we see how we would use a WaitGroup, lets’s make one!
Fortunately, they’re very simple and only require a few atomics. No mutexes or condition variables necessary! (Hence why I call it a primitive).
Here’s the code in TypeScript:
class WaitGroup {
private memory: Int32Array
private offset: number
constructor(memory: Int32Array, offset: number) {
this.memory = memory
this.offset = offset
}
public add(count: number = 1) {
// Very simple, just an atomic add
Atomics.add(this.memory, this.offset, count)
}
public done() {
// Perform a check if we should notify the waiter
// We'll notify the waiter of we set the value to zero
// Due to how atomics return the *old* value, that means
// for a subtract we send when the old value was 1
if (Atomics.sub(this.memory, this.offset, 1) === 1) {
Atomics.notify(this.memory, this.offset)
}
}
public wait(timeout: number = Infinity) {
let lastTime = Date.now()
// Loop until our counter hits zero
while (true) {
// Load our counter and see if we're zero!
const cur = Atomics.load(this.memory, this.offset);
if (cur == 0) {
return true;
}
// Suspend when we're not zero
if (Atomics.wait(this.memory, this.offset, cur, timeout) === 'timed-out') {
return false
}
// remember to update the timeout value whenever we loop!
if (Number.isFinite(timeout)) {
let curTime = Date.now()
let elapsed = curTime - lastTime
timeout -= elapsed
lastTime = curTime
if (timeout <= 0) {
return false
}
}
}
}
public async waitAsync(timeout: number = Infinity) {
// Same thing as above, but with promises now!
let lastTime = Date.now()
while (true) {
const cur = Atomics.load(this.memory, this.offset);
if (cur == 0) {
return true;
}
// Yay promises!
const {async, value} = (Atomics as any).waitAsync(this.memory, this.offset, cur, timeout)
if (async) {
if (await value === 'timed-out') {
return false
}
} else if (value === 'timed-out') {
return false
} else {
// Always ensure we suspend for at least one micro-tick per cycle
await new Promise(res => res(null))
}
if (Number.isFinite(timeout)) {
let curTime = Date.now()
let elapsed = curTime - lastTime
timeout -= elapsed
lastTime = curTime
if (timeout <= 0) {
return false
}
}
}
}
}The general idea is pretty straightforward:
When we call “add”, we atomically increment the counter
When we call “done”, we atomically decrement the counter
We’ll also signal once we hit zero
When we call “wait”, we just loop until the counter hits zero while guaranteeing a suspend in every iteration
Pretty straightforward stuff.
Unfortunately, due to how the standard operates the main thread will always need to send at least one message to the worker. This is due to the fact that the standard requires us to use message passing to pass in shared memory. So, when I say “remove message passing” I’m talking about removing the messages sent from the worker to the main/parent thread.
We do have to reconstruct the wait group inside the worker simply because classes and class instances cannot be passed with message passing. We could use some very clever overrides of global methods to make the system automated - however we would still be reconstructing it on the other end. For explicitness, I’ll leave the manual reconstructions in-plain sight for this series.