Files
2026-07-13 13:23:39 +08:00

631 lines
21 KiB
TypeScript

import type { LoadSignal, LoadSnapshot } from '@crawlee/core';
import { AutoscaledPool } from '@crawlee/core';
import { sleep } from '@crawlee/utils';
import log from '@apify/log';
describe('AutoscaledPool', () => {
let logLevel: number;
beforeAll(() => {
logLevel = log.getLevel();
log.setLevel(log.LEVELS.ERROR);
});
afterAll(() => {
log.setLevel(logLevel);
});
test('should work with concurrency 1', async () => {
const range = [...Array(10).keys()];
const result: number[] = [];
let isFinished = false;
const runTaskFunction = async () => {
if (range.length === 1) {
isFinished = true;
}
return new Promise((resolve) => {
const item = range.shift()!;
result.push(item);
setTimeout(resolve, 5);
});
};
const pool = new AutoscaledPool({
minConcurrency: 1,
maxConcurrency: 1,
runTaskFunction,
isFinishedFunction: async () => Promise.resolve(isFinished),
isTaskReadyFunction: async () => Promise.resolve(!isFinished),
});
await pool.run();
expect(result).toEqual([...Array(10).keys()]);
});
test('should work with concurrency 10', async () => {
const range = [...Array(100).keys()];
const result: number[] = [];
let isFinished = false;
const runTaskFunction = async () => {
if (range.length === 1) {
isFinished = true;
}
return new Promise((resolve) => {
const item = range.shift()!;
result.push(item);
setTimeout(resolve, 5);
});
};
const pool = new AutoscaledPool({
minConcurrency: 10,
maxConcurrency: 10,
runTaskFunction,
isFinishedFunction: async () => Promise.resolve(isFinished),
isTaskReadyFunction: async () => Promise.resolve(!isFinished),
});
await pool.run();
expect(result).toEqual([...Array(100).keys()]);
});
test('enables setting concurrency', async () => {
const range = [...Array(100).keys()];
const result: number[] = [];
let isFinished = false;
const runTaskFunction = async () => {
if (range.length === 1) {
isFinished = true;
}
return new Promise((resolve) => {
const item = range.shift()!;
result.push(item);
setTimeout(resolve, 5);
});
};
const pool = new AutoscaledPool({
// Test initial concurrency setting
minConcurrency: 3,
maxConcurrency: 13,
desiredConcurrency: 9,
runTaskFunction,
isFinishedFunction: async () => isFinished,
isTaskReadyFunction: async () => !isFinished,
});
expect(pool.minConcurrency).toBe(3);
expect(pool.maxConcurrency).toBe(13);
expect(pool.desiredConcurrency).toBe(9);
const promise = pool.run();
// Test setting concurrency
pool.minConcurrency = 4;
pool.maxConcurrency = 14;
pool.desiredConcurrency = 7;
expect(pool.minConcurrency).toBe(4);
expect(pool.maxConcurrency).toBe(14);
expect(pool.desiredConcurrency).toBe(7);
await promise;
expect(result).toEqual([...Array(100).keys()]);
});
describe('should scale correctly', () => {
class MockSystemStatus {
okNow: boolean;
okLately: boolean;
getCurrentStatus: () => { isSystemIdle: boolean };
getHistoricalStatus: () => { isSystemIdle: boolean };
constructor(okNow: boolean, okLately: boolean) {
this.okNow = okNow;
this.okLately = okLately;
this.getCurrentStatus = () => ({ isSystemIdle: this.okNow });
this.getHistoricalStatus = () => ({ isSystemIdle: this.okLately });
}
}
let pool: AutoscaledPool;
let systemStatus: MockSystemStatus;
const cb = () => {};
beforeEach(() => {
systemStatus = new MockSystemStatus(true, true);
pool = new AutoscaledPool({
minConcurrency: 1,
maxConcurrency: 100,
runTaskFunction: async () => {},
isFinishedFunction: async () => false,
isTaskReadyFunction: async () => true,
});
// @ts-expect-error Mock
pool.systemStatus = systemStatus;
});
test('works with low values', () => {
// @ts-expect-error Calling private method
pool._autoscale(cb);
expect(pool.desiredConcurrency).toBe(2);
// @ts-expect-error Calling private method
pool._autoscale(cb);
expect(pool.desiredConcurrency).toBe(2); // because currentConcurrency is not high enough;
// @ts-expect-error Overwriting readonly private prop
pool._currentConcurrency = 2;
// @ts-expect-error Calling private method
pool._autoscale(cb);
expect(pool.desiredConcurrency).toBe(3);
systemStatus.okNow = false; // this should have no effect
// @ts-expect-error Overwriting readonly private prop
pool._currentConcurrency = 3;
// @ts-expect-error Calling private method
pool._autoscale(cb);
expect(pool.desiredConcurrency).toBe(4);
systemStatus.okLately = false;
// @ts-expect-error Calling private method
pool._autoscale(cb);
expect(pool.desiredConcurrency).toBe(3);
});
test('works with high values', () => {
// Should not scale because current concurrency is too low.
pool.desiredConcurrency = 50;
// @ts-expect-error Overwriting readonly private prop
pool._currentConcurrency = Math.floor(pool.desiredConcurrency * pool.desiredConcurrencyRatio) - 1;
systemStatus.okLately = true;
// @ts-expect-error Calling private method
pool._autoscale(cb);
expect(pool.desiredConcurrency).toBe(50);
// Should scale because we bumped up current concurrency.
// @ts-expect-error Overwriting readonly private prop
pool._currentConcurrency = Math.floor(pool.desiredConcurrency * pool.desiredConcurrencyRatio);
// @ts-expect-error Accessing private prop
let newConcurrency = pool.desiredConcurrency + Math.ceil(pool.desiredConcurrency * pool.scaleUpStepRatio);
// @ts-expect-error Calling private method
pool._autoscale(cb);
expect(pool.desiredConcurrency).toEqual(newConcurrency);
// Should scale down.
systemStatus.okLately = false;
// @ts-expect-error Accessing private prop
newConcurrency = pool.desiredConcurrency - Math.ceil(pool.desiredConcurrency * pool.scaleDownStepRatio);
// @ts-expect-error Calling private method
pool._autoscale(cb);
expect(pool.desiredConcurrency).toEqual(newConcurrency);
});
test('works at minConcurrency when currently overloaded', async () => {
let limit = 5;
let concurrencyLog: number[] = [];
let count = 0;
// @ts-expect-error Overwriting readonly private prop
pool.systemStatus.okNow = false;
// @ts-expect-error Overwriting readonly private prop
pool.runTaskFunction = async () => {
await sleep(10);
count++;
};
// @ts-expect-error Overwriting readonly private prop
pool.isFinishedFunction = async () => count >= limit;
// @ts-expect-error Overwriting readonly private prop
pool.isTaskReadyFunction = async () => count < limit;
pool.desiredConcurrency = 10;
// @ts-expect-error Overwriting readonly private prop
pool._currentConcurrency = pool.currentConcurrency;
Object.defineProperty(pool, 'currentConcurrency', {
get() {
return this._currentConcurrency;
},
set(v) {
concurrencyLog.push(v);
this._currentConcurrency = v;
},
});
expect(pool.currentConcurrency).toBe(0);
await pool.run();
expect(concurrencyLog.some((n) => n > 1)).toBe(false);
limit = 50;
concurrencyLog = [];
count = 0;
pool.minConcurrency = 5;
await pool.run();
expect(concurrencyLog.some((n) => n > 5)).toBe(false);
});
});
describe('should throw', () => {
// Turn off unnecessary error logging.
let originalLevel: number;
beforeAll(() => {
originalLevel = log.getLevel();
log.setLevel(log.LEVELS.OFF);
});
afterAll(() => {
log.setLevel(originalLevel);
});
test('when some of the promises throws', async () => {
let counter = 0;
const runTaskFunction = async () => {
counter++;
await sleep(1);
if (counter > 20) throw new Error('some-promise-error');
};
const pool = new AutoscaledPool({
maxConcurrency: 5,
minConcurrency: 5,
runTaskFunction,
isFinishedFunction: async () => counter > 200,
isTaskReadyFunction: async () => true,
});
await expect(pool.run()).rejects.toThrow('some-promise-error');
});
test('when runTaskFunction throws', async () => {
const runTaskFunction = async () => {
await sleep(3);
throw new Error('some-runtask-error');
};
const pool = new AutoscaledPool({
maxConcurrency: 1,
runTaskFunction,
isFinishedFunction: async () => false,
isTaskReadyFunction: async () => true,
});
await expect(pool.run()).rejects.toThrow('some-runtask-error');
});
test('when isFinishedFunction throws', async () => {
let count = 0;
const pool = new AutoscaledPool({
maxConcurrency: 1,
runTaskFunction: async () => {
count++;
},
isFinishedFunction: async () => {
throw new Error('some-finished-error');
},
isTaskReadyFunction: async () => {
return count < 1;
},
});
await expect(pool.run()).rejects.toThrow('some-finished-error');
});
test('when isTaskReadyFunction throws', async () => {
let count = 0;
const pool = new AutoscaledPool({
maxConcurrency: 1,
runTaskFunction: async () => {
count++;
},
isFinishedFunction: async () => false,
isTaskReadyFunction: async () => {
if (count > 1) throw new Error('some-ready-error');
else return true;
},
});
await expect(pool.run()).rejects.toThrow('some-ready-error');
});
});
test('should not handle tasks added after isFinishedFunction returned true', async () => {
const isFinished = async () => count > 10;
let count = 0;
// Run the pool and close it after 3s.
const pool = new AutoscaledPool({
minConcurrency: 3,
runTaskFunction: async () =>
sleep(1).then(() => {
count++;
}),
isFinishedFunction: isFinished,
isTaskReadyFunction: async () => !(await isFinished()),
});
// @ts-expect-error Overwriting readonly private prop
pool.maybeRunIntervalMillis = 5;
await pool.run();
await sleep(10);
expect(count).toBeGreaterThanOrEqual(11);
// Check finished tasks.
expect(count).toBeLessThanOrEqual(13);
});
test('should break and resume when the task queue is empty for a while', async () => {
const finished: number[] = [];
let isFinished = false;
let isTaskReady = true;
let counter = 0;
const pool = new AutoscaledPool({
maxConcurrency: 1,
runTaskFunction: async () => {
await sleep(1);
if (counter === 10) {
isTaskReady = false;
setTimeout(() => {
isTaskReady = true;
}, 10);
}
if (counter === 19) {
isTaskReady = false;
isFinished = true;
}
counter++;
finished.push(Date.now());
},
isFinishedFunction: async () => isFinished,
isTaskReadyFunction: async () => !isFinished && isTaskReady,
});
// @ts-expect-error Overwriting readonly private prop
pool.maybeRunIntervalMillis = 1;
await pool.run();
// Check finished tasks.
expect(finished).toHaveLength(20);
expect(finished[11] - finished[10]).toBeGreaterThan(9);
});
test('should work with loggingIntervalSecs = null', async () => {
const pool = new AutoscaledPool({
minConcurrency: 1,
maxConcurrency: 100,
runTaskFunction: async () => Promise.resolve(),
isFinishedFunction: async () => Promise.resolve(false),
isTaskReadyFunction: async () => Promise.resolve(true),
loggingIntervalSecs: null,
});
// @ts-expect-error Calling private method
pool._autoscale(() => {});
expect(pool.desiredConcurrency).toBe(2);
});
test('should abort', async () => {
let finished = false;
let aborted = false;
const pool = new AutoscaledPool({
runTaskFunction: async () => {
if (!aborted) {
await pool.abort();
aborted = true;
}
},
isFinishedFunction: async () => {
finished = true;
return true;
},
isTaskReadyFunction: async () => !aborted,
});
await pool.run();
expect(finished).toBe(false);
});
test('should only finish after tasks complete', async () => {
let started = false;
let completed = false;
const pool = new AutoscaledPool({
runTaskFunction: async () => {
started = true;
await sleep(100);
completed = true;
},
isFinishedFunction: async () => {
return started;
},
isTaskReadyFunction: async () => {
return !started;
},
});
await pool.run();
expect(started).toBe(true);
expect(completed).toBe(true);
});
test('should pause and resume', async () => {
let count = 0;
const results: number[] = [];
let pauseResolve: (value: unknown) => void;
const pausePromise = new Promise((res) => {
pauseResolve = res;
});
const pool = new AutoscaledPool({
maybeRunIntervalSecs: 0.01,
minConcurrency: 10,
runTaskFunction: async () => {
results.push(count++);
if (count === 20) {
void pool.pause().then(pauseResolve);
}
},
isFinishedFunction: async () => !(count < 50),
isTaskReadyFunction: async () => count < 50,
});
let finished = false;
const runPromise = pool.run();
void runPromise.then(() => {
finished = true;
});
await pausePromise;
expect(count).toBe(20);
expect(finished).toBe(false);
expect(results).toHaveLength(count);
results.forEach((r, i) => expect(r).toEqual(i));
pool.resume();
await runPromise;
expect(count).toBe(50);
expect(finished).toBe(true);
expect(results).toHaveLength(count);
results.forEach((r, i) => expect(r).toEqual(i));
});
test('should timeout after taskTimeoutSecs', async () => {
const runTaskFunction = async () => {
await sleep(1e3);
return 1;
};
const pool = new AutoscaledPool({
minConcurrency: 1,
maxConcurrency: 1,
runTaskFunction,
taskTimeoutSecs: 0.1,
isFinishedFunction: async () => false,
isTaskReadyFunction: async () => true,
});
const now = Date.now();
await expect(pool.run()).rejects.toThrow('runTaskFunction timed out after 0.1 seconds.');
expect(Date.now() - now).toBeGreaterThanOrEqual(100);
});
test('should not timeout if taskTimeoutSecs === 0', async () => {
let finished = false;
const runTaskFunction = async () => {
await sleep(1e3);
finished = true;
return 1;
};
const pool = new AutoscaledPool({
minConcurrency: 1,
maxConcurrency: 1,
runTaskFunction,
taskTimeoutSecs: 0,
isFinishedFunction: async () => finished,
isTaskReadyFunction: async () => !finished,
});
const now = Date.now();
await expect(pool.run()).resolves.toBeUndefined();
expect(Date.now() - now).toBeGreaterThanOrEqual(1e3);
}, 10e3);
test('should not timeout if taskTimeoutSecs not explicitly set', async () => {
let finished = false;
const runTaskFunction = async () => {
await sleep(1e3);
finished = true;
return 1;
};
const pool = new AutoscaledPool({
minConcurrency: 1,
maxConcurrency: 1,
runTaskFunction,
isFinishedFunction: async () => finished,
isTaskReadyFunction: async () => !finished,
});
const now = Date.now();
await expect(pool.run()).resolves.toBeUndefined();
expect(Date.now() - now).toBeGreaterThanOrEqual(1e3);
}, 10e3);
describe('custom load signals', () => {
/** Creates a minimal fake LoadSignal with static pre-seeded snapshots. */
function createFakeLoadSignal(name: string, { overloadedRatio = 0.3, isOverloaded = false } = {}): LoadSignal {
const now = Date.now();
const snapshots: LoadSnapshot[] = Array.from({ length: 5 }, (_, i) => ({
createdAt: new Date(now - (5 - i) * 100),
isOverloaded,
}));
return {
name,
overloadedRatio,
async start() {},
async stop() {},
getSample(sampleDurationMillis?: number) {
if (!sampleDurationMillis) return snapshots;
const cutoff = Date.now() - sampleDurationMillis;
return snapshots.filter((s) => +s.createdAt >= cutoff);
},
};
}
test('overloaded signal prevents concurrency from scaling up', async () => {
const signal = createFakeLoadSignal('proxyHealth', { isOverloaded: true });
let count = 0;
const pool = new AutoscaledPool({
minConcurrency: 1,
maxConcurrency: 10,
runTaskFunction: async () => {
count++;
await sleep(10);
},
isFinishedFunction: async () => count >= 50,
isTaskReadyFunction: async () => count < 50,
systemStatusOptions: { loadSignals: [signal] },
});
await pool.run();
expect(pool.desiredConcurrency).toBe(1);
});
test('signal info appears in SystemStatus.getCurrentStatus()', async () => {
const signal = createFakeLoadSignal('navTimeout', { overloadedRatio: 0.2, isOverloaded: true });
let count = 0;
const pool = new AutoscaledPool({
minConcurrency: 1,
maxConcurrency: 10,
runTaskFunction: async () => {
count++;
await sleep(10);
},
isFinishedFunction: async () => count >= 10,
isTaskReadyFunction: async () => count < 10,
systemStatusOptions: { loadSignals: [signal] },
});
// @ts-expect-error Accessing private prop
const status = pool.systemStatus.getCurrentStatus();
expect(status.loadSignalInfo?.navTimeout?.isOverloaded).toBe(true);
await pool.run();
});
});
});