swaf/src/Throttler.ts

import {TooManyRequestsHttpError} from "./HttpError.js";
import {logger} from "./Logger.js";

export default class Throttler {
    private static readonly throttles: Record<string, Throttle | undefined> = {};

    /**
     * Throttle function; will throw a TooManyRequestsHttpError when the threshold is reached.
     *
     * This throttle is adaptive: it will slowly decrease (linear) until it reaches 0 after {@param resetPeriod} ms.
     * Threshold will hold for {@param holdPeriod} ms.
     *
     * @param action a unique action name (can be used multiple times, but it'll account for a single action).
     * @param max how many times this action can be triggered per id.
     * @param resetPeriod after how much time in ms the throttle will reach 0.
     * @param id an identifier of who triggered the action.
     * @param holdPeriod time in ms after each call before the threshold begins to decrease.
     * @param jailPeriod time in ms for which the throttle will throw when it is triggered.
     */
    public static throttle(
        action: string,
        max: number,
        resetPeriod: number,
        id: string,
        holdPeriod: number = 100,
        jailPeriod: number = 30 * 1000,
    ): void {
        let throttle = this.throttles[action];
        if (!throttle)
            throttle = this.throttles[action] = new Throttle(action, max, resetPeriod, holdPeriod, jailPeriod);
        throttle.trigger(id);
    }

    private constructor() {
        // Disable constructor
    }
}

class Throttle {
    private readonly jailName: string;
    private readonly max: number;
    private readonly resetPeriod: number;
    private readonly holdPeriod: number;
    private readonly jailPeriod: number;
    private readonly triggers: Record<string, {
        count: number,
        lastTrigger?: number,
        jailed?: number;
    } | undefined> = {};

    public constructor(jailName: string, max: number, resetPeriod: number, holdPeriod: number, jailPeriod: number) {
        this.jailName = jailName;
        this.max = max;
        this.resetPeriod = resetPeriod;
        this.holdPeriod = holdPeriod;
        this.jailPeriod = jailPeriod;
    }


    public trigger(id: string) {
        let trigger = this.triggers[id];
        if (!trigger) trigger = this.triggers[id] = {count: 0};

        const currentDate = new Date().getTime();

        if (trigger.jailed && currentDate - trigger.jailed < this.jailPeriod)
            return this.throw(trigger.jailed + this.jailPeriod - currentDate);

        if (trigger.lastTrigger) {
            let timeDiff = currentDate - trigger.lastTrigger;
            if (timeDiff > this.holdPeriod) {
                timeDiff -= this.holdPeriod;
                trigger.count = Math.floor(Math.min(trigger.count, (this.max + 1) * (1 - timeDiff / this.resetPeriod)));
            }
        }

        trigger.count++;
        trigger.lastTrigger = currentDate;

        if (trigger.count > this.max) {
            trigger.jailed = currentDate;

            const unjailedIn = trigger.jailed + this.jailPeriod - currentDate;
            logger.info(`Jail ${this.jailName} triggered by ${id} and will be unjailed in ${unjailedIn}ms.`);

            return this.throw(unjailedIn);
        }
    }

    protected throw(unjailedIn: number) {
        throw new TooManyRequestsHttpError(unjailedIn, this.jailName);
    }
}