Implement semaphores

Co-authored-by: Mark Wubben <[email protected]>

Author: ninevra

Diff for: README.md

@@ -81,3 +81,69 @@ const reserved = await context.reserve(1, 2, 3);
 // `reserved` will be an array containing those values that could be reserved.
 // It could be empty.
 ```
+
+### Semaphores
+
+You can create a [counting semaphore](https://www.guru99.com/semaphore-in-operating-system.html) within a shared context:
+
+```js
+const initialValue = 3; // Must be a non-negative integer.
+const semaphore = context.createSemaphore('my-semaphore', initialValue);
+```
+
+Within the same context, semaphores with the same ID must be created with the same initial value. Semaphores created with a different value are unusable. Their methods will reject with a `SemaphoreCreationError`.
+
+Semaphores have two methods: `acquire()` and `acquireNow()`. Use `acquire()` to decrement the semaphore's value. If the semaphore's value would become negative, instead `acquire()` waits until the semaphore's value is high enough.
+
+```js
+const semaphore = context.createSemaphore('my-semaphore', 3);
+const release = await semaphore.acquire();
+```
+
+`acquire()` returns a function, `release()`, which increments the semaphore's value by the same amount as was acquired.
+
+The semaphore is _managed_: if you don't call `release()`, it'll be run automatically when the test worker exits. Any pending `acquire()` calls will also be removed from the queue at this time.
+
+`acquireNow()` works like `acquire()`, except that if the semaphore can't be decremented immediately, `acquireNow()` rejects with a `SemaphoreDownError` rather than wait.
+
+Semaphores are _weighted_. `acquire()` and `acquireNow()` accept a non-negative integer amount, defaulting to `1`, by which to decrement or increment the value:
+
+```js
+await semaphore.acquire(0);
+await semaphore.acquireNow(2);
+```
+
+You can also pass an amount to `release()` to release just part of the acquisition at a time:
+
+```js
+const release = await semaphore.acquire(3); // Decrements the semaphore by 3
+release(1); // Increments the semaphore by 1
+release(); // Increments the semaphore by the remaining 2
+```
+
+`acquire()` calls resolve in FIFO order. If the current value is `1`, and a call tries to acquire `2`, subsequent `acquire()` calls have to wait, even if they want to acquire just `1`.
+
+`acquireNow()` skips the queue and decrements immediately if possible.
+
+#### Lower-level, unmanaged semaphores
+
+You can create a lower-level, _unmanaged_ semaphore which doesn't have any auto-release behavior. Instead you need to increment the semaphore in code.
+
+```js
+const initialValue = 3; // Must be a non-negative integer.
+const semaphore = context.createUnmanagedSemaphore('my-semaphore', initialValue);
+```
+
+Unmanaged semaphores mustn't use the same ID as a managed semaphore, within the same context. Semaphores with the same ID must be created with the same initial value. Mismatched managed and unmanaged semaphores, or those created with different values are unusable. Their methods will reject with a `SemaphoreCreationError`.
+
+Unmanaged semaphores have three methods. `down()` and `downNow()` decrement the value and `up()` increments:
+
+```js
+await semaphore.down(0);
+await semaphore.downNow(2);
+await semaphore.up(); // `amount` defaults to 1
+```
+
+Like the `acquire()` and `acquireNow()` methods of managed semaphores, `down()` waits for the semaphore's value to be at least the requested amount, while `downNow()` rejects with `SemaphoreDownError` if the value cannot be decremented immediately.
+
+These unmanaged semaphores do not release the "acquired" amount when a test worker exits.

Diff for: source/index.ts

@@ -1,8 +1,10 @@
 import path from 'path';
-import {registerSharedWorker} from 'ava/plugin';
+import {registerSharedWorker, SharedWorker} from 'ava/plugin';
 import never from 'never';
 
-import {Data, MessageType} from './types';
+import {Data, MessageType, SemaphoreCreationFailed} from './types';
+
+type ReceivedMessage = SharedWorker.Plugin.Experimental.ReceivedMessage<Data>;
 
 const protocol = registerSharedWorker<Data>({
 	filename: path.join(__dirname, 'worker.js'),
@@ -17,7 +19,7 @@ export class Lock {
 	}
 
 	async acquire(): Promise<() => void> {
-		// Allow reserve() to be called before the shared worker is availabe.
+		// Allow acquire() to be called before the shared worker is availabe.
 		await protocol.available;
 
 		const message = protocol.publish({
@@ -35,7 +37,6 @@ export class Lock {
 			}
 		}
 
-		/* c8 ignore next 2 */
 		// The above loop will never actually break if the lock is not acquired.
 		return never();
 	}
@@ -63,7 +64,6 @@ export class Lock {
 			}
 		}
 
-		/* c8 ignore next 2 */
 		// The above loop will never actually break if the lock is not acquired.
 		return never();
 	}
@@ -79,13 +79,214 @@ export class LockAcquisitionError extends Error {
 	}
 }
 
+export class ManagedSemaphore {
+	readonly #context: SharedContext;
+
+	constructor(
+		context: SharedContext,
+		public readonly id: string,
+		public readonly initialValue: number
+	) {
+		if (initialValue < 0 || !Number.isSafeInteger(initialValue)) {
+			throw new RangeError('initialValue must be a non-negative safe integer');
+		}
+
+		this.#context = context;
+	}
+
+	async acquire(amount = 1) {
+		if (amount < 0 || !Number.isSafeInteger(amount)) {
+			throw new RangeError('amount must be a non-negative safe integer');
+		}
+
+		// Allow acquire() to be called before the shared worker is availabe.
+		await protocol.available;
+
+		const reply = await downSemaphore(this, this.#context.id, amount, true);
+		return (release = amount) => {
+			if (release < 0 || !Number.isSafeInteger(release) || release > amount) {
+				throw new RangeError('Amount to release must be a non-negative safe integer and <= remaining amount');
+			}
+
+			amount -= release;
+			reply.reply({
+				type: MessageType.SEMAPHORE_RELEASE,
+				amount: release
+			});
+		};
+	}
+
+	async acquireNow(amount = 1) {
+		if (amount < 0 || !Number.isSafeInteger(amount)) {
+			throw new RangeError('amount must be a non-negative safe integer');
+		}
+
+		// Down immediately, which will fail if the protocol is not available.
+		// "Now" should not mean "wait until we're ready."
+
+		const reply = await downSemaphore(this, this.#context.id, amount, false);
+		return (release = amount) => {
+			if (release < 0 || release > amount) {
+				throw new RangeError('Amount to release must be >= 0 and <= remaining amount');
+			}
+
+			amount -= release;
+			reply.reply({
+				type: MessageType.SEMAPHORE_RELEASE,
+				amount: release
+			});
+		};
+	}
+}
+
+export class UnmanagedSemaphore {
+	readonly #context: SharedContext;
+
+	constructor(
+		context: SharedContext,
+		public readonly id: string,
+		public readonly initialValue: number
+	) {
+		if (initialValue < 0 || !Number.isSafeInteger(initialValue)) {
+			throw new RangeError('initialValue must be a non-negative safe integer');
+		}
+
+		this.#context = context;
+	}
+
+	async down(amount = 1) {
+		if (amount < 0 || !Number.isSafeInteger(amount)) {
+			throw new RangeError('amount must be a non-negative safe integer');
+		}
+
+		// Allow down() to be called before the shared worker is availabe.
+		await protocol.available;
+
+		await downSemaphore(this, this.#context.id, amount, true);
+	}
+
+	async downNow(amount = 1) {
+		if (amount < 0 || !Number.isSafeInteger(amount)) {
+			throw new RangeError('amount must be a non-negative safe integer');
+		}
+
+		// Down immediately, which will fail if the protocol is not available.
+		// "Now" should not mean "wait until we're ready."
+
+		await downSemaphore(this, this.#context.id, amount, false);
+	}
+
+	async up(amount = 1) {
+		if (amount < 0 || !Number.isSafeInteger(amount)) {
+			throw new RangeError('amount must be a non-negative safe integer');
+		}
+
+		// Allow up() to be called before the shared worker is availabe.
+		await protocol.available;
+
+		const {id, initialValue} = this;
+		const message = protocol.publish({
+			type: MessageType.SEMAPHORE_UP,
+			contextId: this.#context.id,
+			semaphore: {managed: false, id, initialValue},
+			amount
+		});
+
+		for await (const reply of message.replies()) {
+			if (reply.data.type === MessageType.SEMAPHORE_SUCCEEDED) {
+				return;
+			}
+
+			if (reply.data.type === MessageType.SEMAPHORE_CREATION_FAILED) {
+				throw new SemaphoreCreationError(this, reply.data);
+			}
+		}
+
+		// The above loop will never actually break if the resources are not acquired.
+		return never();
+	}
+}
+
+type Semaphore = ManagedSemaphore | UnmanagedSemaphore;
+
+async function downSemaphore(semaphore: Semaphore, contextId: string, amount: number, wait: boolean): Promise<ReceivedMessage> {
+	const {id, initialValue} = semaphore;
+	const message = protocol.publish({
+		type: MessageType.SEMAPHORE_DOWN,
+		contextId,
+		semaphore: {managed: semaphore instanceof ManagedSemaphore, id, initialValue},
+		amount,
+		wait
+	});
+
+	for await (const reply of message.replies()) {
+		if (reply.data.type === MessageType.SEMAPHORE_SUCCEEDED) {
+			return reply;
+		}
+
+		if (reply.data.type === MessageType.SEMAPHORE_FAILED) {
+			throw new SemaphoreDownError(id, amount);
+		}
+
+		if (reply.data.type === MessageType.SEMAPHORE_CREATION_FAILED) {
+			throw new SemaphoreCreationError(semaphore, reply.data);
+		}
+	}
+
+	// The above loop will never actually break if the resources are not acquired.
+	return never();
+}
+
+export class SemaphoreDownError extends Error {
+	get name() {
+		return 'SemaphoreDownError';
+	}
+
+	constructor(public readonly semaphoreId: string, public readonly amount: number) {
+		super(`Could not immediately decrement with ${amount}`);
+	}
+}
+
+export class SemaphoreCreationError extends Error {
+	readonly semaphoreId: string;
+
+	get name() {
+		return 'SemaphoreCreationError';
+	}
+
+	constructor(semaphore: Semaphore, {initialValue, managed}: SemaphoreCreationFailed) {
+		const initialValueSuffix = `initial value ${semaphore.initialValue} (got ${initialValue})`;
+		if (semaphore instanceof ManagedSemaphore) {
+			if (managed) {
+				super(`Failed to create semaphore: expected ${initialValueSuffix}`);
+			} else {
+				super(`Failed to create semaphore: expected unmanaged and ${initialValueSuffix}`);
+			}
+		} else if (managed) {
+			super(`Failed to create unmanaged semaphore: expected managed and ${initialValueSuffix}`);
+		} else {
+			super(`Failed to create unmanaged semaphore: expected ${initialValueSuffix}`);
+		}
+
+		this.semaphoreId = semaphore.id;
+	}
+}
+
 export class SharedContext {
 	constructor(public readonly id: string) {}
 
 	createLock(id: string): Lock {
 		return new Lock(this, id);
 	}
 
+	createSemaphore(id: string, initialValue: number): ManagedSemaphore {
+		return new ManagedSemaphore(this, id, initialValue);
+	}
+
+	createUnmanagedSemaphore(id: string, initialValue: number): UnmanagedSemaphore {
+		return new UnmanagedSemaphore(this, id, initialValue);
+	}
+
 	async reserve<T extends bigint | number | string>(...values: T[]): Promise<T[]> {
 		// Allow reserve() to be called before the shared worker is availabe.
 		await protocol.available;
@@ -102,7 +303,6 @@ export class SharedContext {
 			}
 		}
 
-		/* c8 ignore next 2 */
 		// The above loop will never actually break if the lock is not acquired.
 		return never();
 	}

Diff for: source/types.d.ts

@@ -5,6 +5,12 @@ export const enum MessageType {
 	LOCK_RELEASE = 13,
 	RESERVE = 20,
 	RESERVED_INDEXES = 21,
+	SEMAPHORE_CREATION_FAILED = 30,
+	SEMAPHORE_DOWN = 31,
+	SEMAPHORE_FAILED = 32,
+	SEMAPHORE_RELEASE = 33,
+	SEMAPHORE_SUCCEEDED = 34,
+	SEMAPHORE_UP = 35,
 }
 
 export type Lock = {
@@ -33,6 +39,43 @@ export type ReservedIndexes = {
 	indexes: number[];
 };
 
+type SemaphoreSetup = {
+	id: string;
+	initialValue: number;
+	managed: boolean;
+};
+
+export type SemaphoreDown = {
+	type: MessageType.SEMAPHORE_DOWN;
+	contextId: string;
+	semaphore: SemaphoreSetup;
+	amount: number;
+	wait: boolean;
+};
+
+export type SemaphoreUp = {
+	type: MessageType.SEMAPHORE_UP;
+	contextId: string;
+	semaphore: SemaphoreSetup;
+	amount: number;
+};
+
+export type SemaphoreRelease = {
+	type: MessageType.SEMAPHORE_RELEASE;
+	amount: number;
+};
+
+export type SemaphoreResult = {
+	type: MessageType.SEMAPHORE_SUCCEEDED | MessageType.SEMAPHORE_FAILED;
+};
+
+export type SemaphoreCreationFailed = {
+	type: MessageType.SEMAPHORE_CREATION_FAILED;
+	initialValue: number;
+	managed: boolean;
+};
+
 export type Data =
 	Lock | Locked | LockRelease |
-	Reservation | ReservedIndexes;
+	Reservation | ReservedIndexes |
+	SemaphoreDown | SemaphoreResult | SemaphoreUp | SemaphoreRelease | SemaphoreCreationFailed;