Nodejs implementation of the Golang waitgroup.

waitgroup.ts

import { EventEmitter } from "events"

export class WaitGroup extends EventEmitter {
	public state: number

	constructor(initialState: number = 0) {
		if (initialState < 0) {
			throw new RangeError()
		}
		super()
		this.state = initialState
	}

	public Add(delta: number = 1): number {
		this.state += delta
		if (this.state < 0) {
			throw new RangeError()
		}

		this.emit("change", this.state)
		if (this.state == 0) {
			this.emit("done")
		}

		return this.state
	}

	public Done(): number {
		return this.Add(-1)
	}

	Wait(): Promise<void>
	Wait(cb: () => void): void
	public Wait(cb?: () => void) {
		// Handle when wait group is already drained.
		if (this.state == 0) {
			// Callback style
			if (cb) {
				return cb()
			}

			// Promise style
			return Promise.resolve()
		}

		// Callback style
		if (cb) {
			this.once("done", cb)
			return
		}

		// Promise style
		return new Promise<void>(resolve => {
			this.once("done", resolve)
		})
	}
}