@stackbit/utils/src/task-queue.ts

Stackbit utilities

let taskTagCounter = 0;

class Task {
    readonly tag: string;
    readonly promise: Promise<any>;
    run!: () => Promise<any>;

    constructor(job: () => Promise<any>, tag?: string) {
        this.tag = tag ?? 'task-' + taskTagCounter++;
        this.promise = new Promise((resolve, reject) => {
            this.run = function () {
                return job().then(resolve, reject);
            };
        });
    }
}

export interface TaskQueueOptions {
    limit?: number;
    interval?: number;
    debug?: boolean;
}

export default class TaskQueue {
    private readonly debug: boolean;
    private runCount: number;
    private lastRunTime: [number, number] | null;
    private timeout: NodeJS.Timeout | null;
    private taskQueue: Task[];

    constructor(options?: TaskQueueOptions) {
        const limit = options?.limit ?? null;
        const interval = options?.interval ?? null;
        this.debug = options?.debug ?? false;
        this.runCount = 0;
        this.lastRunTime = null;
        this.timeout = null;
        this.taskQueue = [];

        if (interval) {
            this._runTask = this._runTaskInterval.bind(this, this._runTask.bind(this), interval);
        }
        if (limit) {
            this._runTask = this._runTaskLimit.bind(this, this._runTask.bind(this), limit);
        }
    }

    addTask<T>(job: () => Promise<T>, tag?: string): Promise<T> {
        const task = new Task(job, tag);
        this.taskQueue.push(task);
        if (this.debug) {
            console.log(`[TaskQueue] submitted task to queue, task: ${task.tag}, queue size: ${this.taskQueue.length}, running tasks: ${this.runCount}`);
        }
        this._executeNextTask();
        return task.promise;
    }

    clearQueue() {
        this.taskQueue = [];
        if (this.timeout) {
            clearTimeout(this.timeout);
            this.timeout = null;
        }
    }

    private _executeNextTask() {
        if (this.taskQueue.length > 0) {
            this._runTask();
        }
    }

    private _runTaskLimit(runTask: () => void, limit: number) {
        if (this.runCount < limit) {
            runTask();
        }
        if (this.debug) {
            console.log(
                `[TaskQueue] task run count limit (${limit}) reached, queue size: ${this.taskQueue.length}, ` +
                    `running tasks: ${this.runCount}, waiting for previous tasks to finish`
            );
        }
    }

    private _runTaskInterval(runTask: () => void, interval: number) {
        if (!this.lastRunTime) {
            // if this is a first task, run it immediately
            runTask();
        } else if (!this.timeout) {
            // if there is no timeout, check if the time from last run is greater
            // than the allowed interval and run the task, otherwise set a timeout
            // to run the task at the right interval
            const diff = process.hrtime(this.lastRunTime);
            const diffMs = diff[0] * 1000 + diff[1] / 1000000;
            if (diffMs >= interval) {
                runTask();
            } else {
                this.timeout = setTimeout(runTask, interval - diffMs);
                if (this.debug) {
                    console.log(
                        `[TaskQueue] task interval is less than allowed (${interval}) reached, queue size: ` +
                            `${this.taskQueue.length}, running tasks: ${this.runCount}, waiting for ${interval - diffMs}ms`
                    );
                }
            }
        }
        // If the timeout already set, then the next call to _runTask will call
        // _executeNextTask at the end and will set a new timeout for the
        // following task.
    }

    private _runTask() {
        // _runTask is called from _executeNextTask only when the task queue is
        // not empty, and timeout is set only when there is at least one task.
        const task = this.taskQueue.shift()!;
        if (this.timeout) {
            clearTimeout(this.timeout);
            this.timeout = null;
        }
        this.runCount += 1;
        this.lastRunTime = process.hrtime();
        if (this.debug) {
            console.log(`[TaskQueue] running task: ${task.tag}, queue size: ${this.taskQueue.length}, running tasks: ${this.runCount}`);
        }
        task.run().then(() => {
            this.runCount -= 1;
            if (this.debug) {
                console.log(`[TaskQueue] finished running task: ${task.tag}, queue size: ${this.taskQueue.length}, running tasks: ${this.runCount}`);
            }
            this._executeNextTask();
        });
        this._executeNextTask();
    }
}