zero-overhead-keyed-promise-lock
Version:
An efficient keyed Promise lock for Node.js projects, ensuring the mutually exclusive execution of tasks associated with the same key. Key features include active key metrics and the ability to gracefully await the completion of all currently executing or
352 lines (305 loc) • 14 kB
text/typescript
/**
* Copyright 2025 Ori Cohen https://github.com/ori88c
* https://github.com/ori88c/zero-overhead-keyed-promise-lock
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
/**
* General note on testing concurrency in unit tests:
* While ideal tests follow a strict Arrange-Act-Assert structure, rigorously testing
* concurrency-oriented components often requires validating intermediate states.
* Incorrect intermediate states can compromise the entire component's correctness,
* making their verification essential.
*
* As with everything in engineering, this comes at a cost: verbosity.
* Given that resilience is the primary goal, this is a small price to pay.
*/
import { ZeroOverheadKeyedLock } from './zero-overhead-keyed-promise-lock';
const sleep = (ms: number) => new Promise((res) => setTimeout(res, ms));
interface MockErrorMessageInfo {
round: number;
key: string;
}
describe('ZeroOverheadKeyedLock tests', () => {
describe('Happy path tests', () => {
test('executeExclusive: should return the expected value when succeeds', async () => {
// Arrange.
const lock = new ZeroOverheadKeyedLock<number>();
const key = 'user-9799789';
const expectedValue = -295;
const task = async (): Promise<number> => {
return expectedValue;
};
// Pre-action validations.
expect(lock.isActiveKey(key)).toBe(false);
expect(lock.getCurrentExecution(key)).toBeUndefined();
expect(lock.activeKeysCount).toBe(0);
expect(lock.activeKeys).toEqual([]);
// Act.
const actualValue = await lock.executeExclusive(key, task);
// Post-action validations.
expect(actualValue).toBe(expectedValue);
expect(lock.isActiveKey(key)).toBe(false);
expect(lock.getCurrentExecution(key)).toBeUndefined();
expect(lock.activeKeysCount).toBe(0);
expect(lock.activeKeys).toEqual([]);
});
test('waitForAllExistingTasksToComplete: should resolve immediately if no task currently executes', async () => {
const lock = new ZeroOverheadKeyedLock<void>();
await lock.waitForAllExistingTasksToComplete();
expect(lock.getCurrentExecution('mock_key')).toBeUndefined();
expect(lock.activeKeysCount).toBe(0);
});
// prettier-ignore
test(
'executeExclusive: should process only one task at a time per key, and waitForAllExistingTasksToComplete ' +
'should resolve only after *all* the currently pending and processed tasks are completed',
async () => {
jest.useFakeTimers();
const lock = new ZeroOverheadKeyedLock<void>();
const numberOfKeys = 176;
const keys: readonly string[] = new Array<number>(numberOfKeys)
.fill(0)
.map((_, i) => i * 20)
.map((num) => `user-${num}`);
const roundToExecuteExclusivePromises: Promise<void>[][] = [];
const initialBackpressurePerKey = 15;
const taskDurationMs = 8000;
let actualCompletedTasksCount = 0;
const createTask = async (): Promise<void> => {
await sleep(taskDurationMs);
++actualCompletedTasksCount;
};
const validateKeysActivity = (shouldBePresent: boolean): void => {
for (const key of keys) {
expect(lock.isActiveKey(key)).toBe(shouldBePresent);
const currentExecutionPromise = lock.getCurrentExecution(key);
if (shouldBePresent) {
expect(currentExecutionPromise).toBeDefined();
} else {
expect(currentExecutionPromise).toBeUndefined();
}
}
};
// Initial validation.
expect(lock.activeKeysCount).toBe(0);
expect(lock.activeKeys).toEqual([]);
validateKeysActivity(false);
// Create a burst of tasks, inducing backpressure on the lock.
// Add `initialBackpressurePerKey` pending tasks per key.
for (let round = 0; round < initialBackpressurePerKey; ++round) {
const pendingExecuteExclusivePromises: Promise<void>[] = keys.map((key) =>
lock.executeExclusive(key, createTask),
);
roundToExecuteExclusivePromises.push(pendingExecuteExclusivePromises);
// Trigger an event-loop iteration.
await jest.advanceTimersByTimeAsync(0);
validateKeysActivity(true);
expect(lock.activeKeysCount).toBe(keys.length);
expect(lock.activeKeys).toEqual(keys);
}
// Currrently, the number of pending tasks is initialBackpressurePerKey * numberOfKeys.
let allTasksCompleted = false;
const waitForCompletionOfAllTasksPromise: Promise<void> = (async () => {
await lock.waitForAllExistingTasksToComplete();
allTasksCompleted = true;
})();
let expectedCompletedTasksCount = 0;
for (let round = 0; round < initialBackpressurePerKey; ++round) {
expect(allTasksCompleted).toBe(false);
expect(lock.activeKeysCount).toBe(keys.length);
expect(lock.activeKeys).toEqual(keys);
validateKeysActivity(true);
// Simulate the completion of one task per key. This occurs because all tasks
// have the same duration, and each key has an equal number of remaining tasks.
await Promise.race([
waitForCompletionOfAllTasksPromise,
jest.advanceTimersByTimeAsync(taskDurationMs),
]);
await Promise.all(roundToExecuteExclusivePromises[round]);
// Each round, we complete one pending task per key.
expectedCompletedTasksCount += numberOfKeys;
expect(actualCompletedTasksCount).toBe(expectedCompletedTasksCount);
}
await waitForCompletionOfAllTasksPromise;
expect(lock.activeKeysCount).toBe(0);
expect(lock.activeKeys).toEqual([]);
expect(allTasksCompleted).toBe(true);
validateKeysActivity(false);
jest.useRealTimers();
},
);
test('getCurrentExecution should return the active task promise during execution', async () => {
type ResultType = Record<string, number>;
const mockResult: ResultType = { a: 1, b: 2 };
let resolveTask: (value?: unknown) => void;
const lock = new ZeroOverheadKeyedLock<ResultType>();
const key = 'mock-key';
const taskPromise = new Promise<ResultType>((res) => (resolveTask = res));
const executeExclusivePromise = lock.executeExclusive(key, () => taskPromise);
const validationRounds = 24;
let currentExecution: Promise<ResultType>;
for (let attempt = 0; attempt < validationRounds; ++attempt) {
expect(lock.isActiveKey(key)).toBe(true);
// getCurrentExecution should return the same Promise instance as long as
// the current task is ongoing.
if (currentExecution) {
expect(lock.getCurrentExecution(key)).toBe(currentExecution);
} else {
currentExecution = lock.getCurrentExecution(key);
expect(currentExecution).toBeDefined();
}
expect(lock.activeKeysCount).toBe(1);
expect(lock.activeKeys).toEqual([key]);
}
resolveTask(mockResult);
const result = await executeExclusivePromise;
expect(result).toBe(mockResult);
expect(lock.activeKeysCount).toBe(0);
expect(lock.activeKeys).toEqual([]);
expect(lock.isActiveKey(key)).toBe(false);
expect(lock.getCurrentExecution(key)).toBeUndefined();
});
});
describe('Negative path tests', () => {
test('executeExclusive & getCurrentExecution: should reject when the key is empty or not a string', async () => {
const lock = new ZeroOverheadKeyedLock<string>();
const createTask = () => Promise.resolve<string>('mock-result');
const invalidKeys = [
'',
0 as unknown as string,
4 as unknown as string,
-53.431 as unknown as string,
undefined as string,
null as string,
true as unknown as string,
{ prop1: 'value1' } as unknown as string,
];
for (const key of invalidKeys) {
await expect(() => lock.executeExclusive(key, createTask)).rejects.toThrow();
expect(() => lock.getCurrentExecution(key)).toThrow();
}
expect(lock.activeKeysCount).toBe(0);
expect(lock.activeKeys).toEqual([]);
await lock.waitForAllExistingTasksToComplete();
});
test('executeExclusive: should return the expected error when throws', async () => {
const lock = new ZeroOverheadKeyedLock<string>();
const key = 'mock-key';
const expectedError = new Error('mock error');
const createTask = async (): Promise<string> => {
throw expectedError;
};
expect.assertions(5);
try {
await lock.executeExclusive(key, createTask);
} catch (err) {
expect(err).toBe(expectedError);
// The event-driven eviction mechanism should remove the key after the task completes,
// as long as no other task is pending (i.e., no backpressure).
expect(lock.isActiveKey(key)).toBe(false);
expect(lock.getCurrentExecution(key)).toBeUndefined();
}
expect(lock.activeKeysCount).toBe(0);
expect(lock.activeKeys).toEqual([]);
await lock.waitForAllExistingTasksToComplete();
});
// prettier-ignore
test(
'executeExclusive: should process only one task at a time per key, and waitForAllExistingTasksToComplete ' +
'should resolve only after *all* the currently pending and processed tasks are completed, ' +
'with all tasks rejecting',
async () => {
jest.useFakeTimers();
const lock = new ZeroOverheadKeyedLock<void>();
const numberOfKeys = 96;
const keys: readonly string[] = new Array<number>(numberOfKeys)
.fill(0)
.map((_, i) => i * 25)
.map((num) => `user-${num}`);
const roundToExecuteExclusivePromises: Promise<void>[][] = [];
const initialBackpressurePerKey = 12;
const taskDurationMs = 6000;
const createErrorMessage = (info: MockErrorMessageInfo): string =>
`Task for ${info.key} failed during round ${info.round}`;
const task = async (info: MockErrorMessageInfo): Promise<void> => {
await sleep(taskDurationMs);
throw new Error(createErrorMessage(info));
};
const validateKeysActivity = (shouldBePresent: boolean): void => {
for (const key of keys) {
expect(lock.isActiveKey(key)).toBe(shouldBePresent);
const currentExecutionPromise = lock.getCurrentExecution(key);
if (shouldBePresent) {
expect(currentExecutionPromise).toBeDefined();
} else {
expect(currentExecutionPromise).toBeUndefined();
}
}
};
// Initial validation.
expect(lock.activeKeysCount).toBe(0);
expect(lock.activeKeys).toEqual([]);
validateKeysActivity(false);
// Create a burst of tasks, inducing backpressure on the lock.
// Add `initialBackpressurePerKey` pending tasks per key.
for (let round = 0; round < initialBackpressurePerKey; ++round) {
const pendingExecuteExclusivePromises: Promise<void>[] = keys.map((key) =>
lock.executeExclusive(key, () => task({ round, key })),
);
roundToExecuteExclusivePromises.push(pendingExecuteExclusivePromises);
// Trigger an event-loop iteration.
await jest.advanceTimersByTimeAsync(0);
validateKeysActivity(true);
expect(lock.activeKeysCount).toBe(keys.length);
expect(lock.activeKeys).toEqual(keys);
}
// Currrently, the number of pending tasks is initialBackpressurePerKey * numberOfKeys.
let allTasksCompleted = false;
const waitForCompletionOfAllTasksPromise: Promise<void> = (async () => {
await lock.waitForAllExistingTasksToComplete();
allTasksCompleted = true;
})();
for (let round = 0; round < initialBackpressurePerKey; ++round) {
expect(allTasksCompleted).toBe(false);
expect(lock.activeKeysCount).toBe(keys.length);
expect(lock.activeKeys).toEqual(keys);
validateKeysActivity(true);
// Simulate the completion of one task per key. This occurs because all tasks
// have the same duration, and each key has an equal number of remaining tasks.
await Promise.allSettled([
...roundToExecuteExclusivePromises[round],
jest.advanceTimersByTimeAsync(taskDurationMs),
]);
// Each round, we complete one pending task per key.
let i = 0;
for (const key of keys) {
try {
await roundToExecuteExclusivePromises[round][i++];
expect(true).toBe(false); // The flow should not reach this point.
} catch (err) {
const expectedMessage = createErrorMessage({ round, key });
expect(err.message).toEqual(expectedMessage);
}
}
}
await waitForCompletionOfAllTasksPromise;
expect(lock.activeKeysCount).toBe(0);
expect(lock.activeKeys).toEqual([]);
expect(allTasksCompleted).toBe(true);
validateKeysActivity(false);
jest.useRealTimers();
},
);
});
});