Blame · RateLimiter.ts

b69ab31	1	/**
b69ab31	2	* Copyright (c) Meta Platforms, Inc. and affiliates.
b69ab31	3	*
b69ab31	4	* This source code is licensed under the MIT license found in the
b69ab31	5	* LICENSE file in the root directory of this source tree.
b69ab31	6	*/
b69ab31	7
b69ab31	8	import {TypedEventEmitter} from './TypedEventEmitter';
b69ab31	9
b69ab31	10	type Id = number;
b69ab31	11
b69ab31	12	/**
b69ab31	13	* Rate limits requests to run an arbitrary task.
b69ab31	14	* Up to `maxSimultaneousRunning` tasks can run at once,
b69ab31	15	* further requests will be queued and run when a running task finishes.
b69ab31	16	*
b69ab31	17	* Usage:
b69ab31	18	* ```
b69ab31	19	* const rateLimiter = new RateLimiter(5);
b69ab31	20	* const result = await rateLimiter.enqueueRun(() => {
b69ab31	21	* // ...do arbitrary async work...
b69ab31	22	* });
b69ab31	23	* ```
b69ab31	24	*/
b69ab31	25	export class RateLimiter {
b69ab31	26	private queued: Array<Id> = [];
b69ab31	27	private running: Array<Id> = [];
b69ab31	28	private runs = new TypedEventEmitter<'run', Id>();
b69ab31	29
b69ab31	30	constructor(
b69ab31	31	private maxSimultaneousRunning: number,
b69ab31	32	private log?: (s: string) => unknown,
b69ab31	33	) {}
b69ab31	34
b69ab31	35	private nextId = 1;
b69ab31	36	private generateId(): Id {
b69ab31	37	return this.nextId++;
b69ab31	38	}
b69ab31	39
b69ab31	40	async enqueueRun<T>(runner: () => Promise<T>): Promise<T> {
b69ab31	41	const id = this.generateId();
b69ab31	42
b69ab31	43	this.queued.push(id);
b69ab31	44	this.tryDequeueNext();
b69ab31	45
b69ab31	46	if (!this.running.includes(id)) {
b69ab31	47	this.log?.(`${this.running.length} tasks are already running, enqueuing ID:${id}`);
b69ab31	48	await new Promise(res => {
b69ab31	49	this.runs.on('run', ran => {
b69ab31	50	if (ran === id) {
b69ab31	51	this.log?.(`now allowing ID:${id} to run`);
b69ab31	52	res(undefined);
b69ab31	53	}
b69ab31	54	});
b69ab31	55	});
b69ab31	56	}
b69ab31	57
b69ab31	58	try {
b69ab31	59	return await runner();
b69ab31	60	} finally {
b69ab31	61	this.notifyFinished(id);
b69ab31	62	}
b69ab31	63	}
b69ab31	64
b69ab31	65	private notifyFinished(id: Id): void {
b69ab31	66	this.running = this.running.filter(running => running !== id);
b69ab31	67	this.tryDequeueNext();
b69ab31	68	}
b69ab31	69
b69ab31	70	private tryDequeueNext() {
b69ab31	71	if (this.running.length < this.maxSimultaneousRunning) {
b69ab31	72	const toRun = this.queued.shift();
b69ab31	73	if (toRun != null) {
b69ab31	74	this.run(toRun);
b69ab31	75	}
b69ab31	76	}
b69ab31	77	}
b69ab31	78
b69ab31	79	private run(id: Id) {
b69ab31	80	this.running.push(id);
b69ab31	81	this.runs.emit('run', id);
b69ab31	82	}
b69ab31	83	}