magnitude-test
Version:
A TypeScript client for running automated UI tests through the Magnitude testing platform
94 lines (93 loc) • 4.4 kB
JavaScript
// packages/magnitude-test/src/runner/workerPool.ts
/**
* A simple worker pool to run async tasks with a concurrency limit and support for early abortion.
*/
export class WorkerPool {
concurrency;
/**
* Creates an instance of WorkerPool.
* @param concurrency The maximum number of tasks to run concurrently. Must be at least 1.
*/
constructor(concurrency) {
this.concurrency = Math.max(1, concurrency);
}
/**
* Runs the given asynchronous tasks with the specified concurrency.
*
* @template T The type of the result returned by each task.
* @param tasks An array of functions, each returning a Promise<T>. Each function receives an AbortSignal.
* @param checkResultForAbort An optional function that checks the result of a completed task. If it returns true, the pool will abort further processing.
* @returns A Promise resolving to a WorkerPoolResult<T> object.
*/
async runTasks(tasks, checkResultForAbort = () => false) {
const abortController = new AbortController();
const { signal } = abortController;
const taskQueue = tasks.map((task, index) => ({ task, index }));
const results = new Array(tasks.length).fill(undefined);
const runningWorkers = new Set();
const runWorker = async () => {
while (taskQueue.length > 0) {
if (signal.aborted) {
break; // Stop processing if aborted
}
const taskItem = taskQueue.shift();
if (!taskItem)
continue; // Should not happen if queue.length > 0
const { task, index } = taskItem;
try {
// Check signal *before* starting the potentially long task
if (signal.aborted) {
// Task skipped due to prior abort
continue;
}
const result = await task(signal);
results[index] = result;
// Check if this result triggers an abort, only if not already aborted
if (!signal.aborted && checkResultForAbort(result)) {
abortController.abort();
}
}
catch (error) {
results[index] = undefined; // Mark result as undefined on error
if (!signal.aborted) {
abortController.abort(); // Abort on any task error
}
}
}
};
const startWorkers = () => {
while (runningWorkers.size < this.concurrency && taskQueue.length > 0 && !signal.aborted) {
const workerPromise = runWorker().finally(() => {
runningWorkers.delete(workerPromise);
// If not aborted and there are still tasks, try starting another worker
// This ensures we maintain concurrency level when a worker finishes
if (!signal.aborted && taskQueue.length > 0) {
startWorkers();
}
});
runningWorkers.add(workerPromise);
}
};
startWorkers(); // Kick off the initial workers
// Wait for all active workers to complete their current task or stop due to abort
// We need to wait until the set is empty, indicating all started workers have finished.
while (runningWorkers.size > 0) {
try {
// Wait for any worker to finish
await Promise.race(runningWorkers);
}
catch (e) {
// Errors within tasks are caught inside runWorker,
// this catch is for unexpected issues with Promise.race or the worker management itself.
console.error("Unexpected error waiting for worker:", e);
if (!signal.aborted) {
abortController.abort();
}
}
}
// Final check: If the queue still has items, it means we aborted early.
const completed = !signal.aborted && taskQueue.length === 0;
// Caller (TestRunner) is responsible for handling UI updates for cancelled tasks based on results array
return { completed, results };
}
}