631 lines
21 KiB
TypeScript
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();
|
|
});
|
|
});
|
|
});
|