From 79a98d235c40b2734e6d739b67b13f67545a353d Mon Sep 17 00:00:00 2001 From: Gilad S Date: Fri, 20 Dec 2024 00:00:19 +0200 Subject: [PATCH] fix(`withLock`): improve queue efficiency --- src/DisposedError.ts | 2 +- src/withLock.ts | 145 +++++++++++++++--------------------------- test/withLock.test.ts | 9 +++ 3 files changed, 61 insertions(+), 95 deletions(-) diff --git a/src/DisposedError.ts b/src/DisposedError.ts index ff22144..d83c498 100644 --- a/src/DisposedError.ts +++ b/src/DisposedError.ts @@ -1,5 +1,5 @@ /** - * This error is throw when an object is disposed and a method is called on it. + * This error is thrown when an object is disposed and a method is called on it. * You can use this error to check if an error is caused by a disposed object. * You can also use this error to throw when an object is disposed and a method is called on it. */ diff --git a/src/withLock.ts b/src/withLock.ts index c338b4d..e97905f 100644 --- a/src/withLock.ts +++ b/src/withLock.ts @@ -1,4 +1,4 @@ -const locks = new Map>>(); +const locks = new Map void)[], onDelete: (() => void)[]]>>(); /** * Only allow one instance of the callback to run at a time for a given `scope` and `key`. @@ -25,44 +25,38 @@ export async function withLock( if (callback == null) throw new Error("callback is required"); - while (locks.get(scope)?.has(key)) { - if (acquireLockSignal?.aborted) - throw acquireLockSignal.reason; + if (acquireLockSignal?.aborted) + throw acquireLockSignal.reason; - try { - if (acquireLockSignal != null) { - const acquireLockPromise = createAbortSignalAbortPromise(acquireLockSignal); - - await Promise.race([ - acquireLockPromise.promise, - locks.get(scope)?.get(key) - ]); - - acquireLockPromise.dispose(); - } else - await locks.get(scope)?.get(key); - } catch (err) { - // we only need to wait here for the promise to resolve, we don't care about the result - } - - if (acquireLockSignal?.aborted) - throw acquireLockSignal.reason; + let keyMap = locks.get(scope); + if (keyMap == null) { + keyMap = new Map(); + locks.set(scope, keyMap); } - const promise = callback.call(scope); - - if (!locks.has(scope)) - locks.set(scope, new Map()); - - locks.get(scope)!.set(key, promise); + let [queue, onDelete] = keyMap.get(key) || []; + if (queue != null && onDelete != null) + await createQueuePromise(queue, acquireLockSignal); + else { + queue = []; + onDelete = []; + keyMap.set(key, [queue, onDelete]); + } try { - return await promise; + return await callback.call(scope); } finally { - locks.get(scope)?.delete(key); + if (queue.length > 0) + queue.shift()!(); + else { + locks.get(scope)?.delete(key); + + if (locks.get(scope)?.size === 0) + locks.delete(scope); - if (locks.get(scope)?.size === 0) - locks.delete(scope); + while (onDelete.length > 0) + onDelete.shift()!(); + } } } @@ -108,43 +102,14 @@ export async function acquireLock( * Wait for a lock to be released for a given `scope` and `key`. */ export async function waitForLockRelease(scope: any, key: string, signal?: AbortSignal): Promise { - // eslint-disable-next-line no-constant-condition - while (true) { - if (signal?.aborted) - throw signal.reason; - - try { - if (signal != null) { - const signalPromise = createAbortSignalAbortPromise(signal); - - await Promise.race([ - signalPromise.promise, - locks.get(scope)?.get(key) - ]); - - signalPromise.dispose(); - } else - await locks.get(scope)?.get(key); - } catch (err) { - // we only need to wait here for the promise to resolve, we don't care about the result - } - - if (signal?.aborted) - throw signal.reason; - - if (locks.get(scope)?.has(key)) - continue; - - await Promise.resolve(); // wait for a microtask to run, so other pending locks can be registered - - if (signal?.aborted) - throw signal.reason; - - if (locks.get(scope)?.has(key)) - continue; + if (signal?.aborted) + throw signal.reason; + const [queue, onDelete] = locks.get(scope)?.get(key) ?? []; + if (queue == null || onDelete == null) return; - } + + await createQueuePromise(onDelete, signal); } export type Lock = { @@ -154,36 +119,28 @@ export type Lock = { [Symbol.dispose](): void }; -function createControlledPromise() { - let resolve: (value: T | Promise) => void; - let reject: (reason?: any) => void; +function createQueuePromise(queue: (() => void)[], signal?: AbortSignal) { + if (signal == null) + return new Promise((accept) => void queue.push(accept)); - const promise = new Promise((accept, fail) => { - resolve = accept; - reject = fail; - }); + return new Promise((accept, reject) => { + function onAcquireLock() { + signal!.removeEventListener("abort", onAbort); + accept(); + } - return { - promise, - resolve: resolve!, - reject: reject! - }; -} + const queueLength = queue.length; -function createAbortSignalAbortPromise(signal: AbortSignal) { - const acquireLockPromise = createControlledPromise(); + function onAbort() { + const itemIndex = queue.lastIndexOf(onAcquireLock, queueLength); + if (itemIndex >= 0) + queue.splice(itemIndex, 1); - const onAbort = () => { - acquireLockPromise.resolve(); - signal.removeEventListener("abort", onAbort); - }; - signal.addEventListener("abort", onAbort); - - return { - promise: acquireLockPromise.promise, - dispose() { - signal.removeEventListener("abort", onAbort); + signal!.removeEventListener("abort", onAbort); + reject(signal!.reason); } - }; -} + queue.push(onAcquireLock); + signal.addEventListener("abort", onAbort); + }); +} diff --git a/test/withLock.test.ts b/test/withLock.test.ts index a41ded5..946fe4d 100644 --- a/test/withLock.test.ts +++ b/test/withLock.test.ts @@ -183,12 +183,21 @@ describe("withLock", () => { const lock2Controller = new AbortController(); const lock2Promise = acquireLock(scope1, key1, lock2Controller.signal); + const lock3Controller = new AbortController(); + const lock3Promise = acquireLock(scope1, key1, lock3Controller.signal); + lock2Controller.abort(new TestError()); await expect(lock2Promise).rejects.toBeInstanceOf(TestError); lock1.dispose(); await new Promise((accept) => setTimeout(accept, 0)); + const lock3 = await lock3Promise; + expect(isLockActive(scope1, key1)).toBe(true); + + lock3.dispose(); + await new Promise((accept) => setTimeout(accept, 0)); + expect(isLockActive(scope1, key1)).toBe(false); });