TypeScript Retry Budget

budget.ts

import { strict as assert } from "assert";

type BudgetOptions = {
    ttlInSeconds: number,
    minPerSecond: number,
    retryPercent: number,
    now?: () => number,
}

export class RetryBudget {
    private bucket: TokenBucket;
    private depositAmount: number;
    private withdrawAmount: number;

    constructor({ ttlInSeconds, minPerSecond, retryPercent, now = Date.now }: BudgetOptions) {
        assert(ttlInSeconds >= 1 && ttlInSeconds <= 60, "ttl must be between 1 and 60 seconds");
        assert(retryPercent >= 0 && retryPercent <= 1000, "retryPercent must be between 0.0 and 1000");
        assert(minPerSecond >= 0, "minPerSecond must be non-negative");

        let deposit: number, withdraw: number;
        if (retryPercent === 0) {
            deposit = 0;
            withdraw = 1;
        } else if (retryPercent <= 1.0) {
            deposit = 1;
            withdraw = 1.0 / retryPercent;
        } else {
            deposit = 1000;
            withdraw = 1000.0 / retryPercent;
        }

        const reserve = minPerSecond * ttlInSeconds * withdraw;

        this.bucket = new TokenBucket(ttlInSeconds * 1000, reserve, now);
        this.depositAmount = deposit;
        this.withdrawAmount = withdraw;
    }

    deposit() {
        this.bucket.put(this.depositAmount);
    }

    withdraw(): boolean {
        return this.bucket.get(this.withdrawAmount);
    }

    balance(): number {
        return Math.floor(this.bucket.count() / this.withdrawAmount);
    }
}

class TokenBucket {
    private windows: WindowedAdder;
    private reserve: number;

    constructor(millisPerWindow: number, reserve: number, now: () => number) {
        this.windows = new WindowedAdder(millisPerWindow / 10, 10, now);
        this.reserve = reserve;
    }

    put(n: number) {
        this.windows.add(n);
    }

    get(n: number): boolean {
        const ok = this.count() >= n;
        if (ok) {
            this.windows.add(-n);
        }
        return ok;
    }

    count(): number {
        return this.windows.sum() + this.reserve;
    }
}

class WindowedAdder {
    private millisPerWindow: number;
    private windows: number[];
    private generation: number;
    private time: number;
    private now: () => number;

    constructor(millis: number, windows: number, now: () => number) {
        this.millisPerWindow = millis;
        this.windows = new Array(windows);
        this.windows.fill(0);
        this.generation = 0;
        this.time = now();
        this.now = now;
    }

    add(x: number) {
        this.expire();
        this.windows[this.generation] += x;
    }

    sum(): number {
        this.expire();
        return this.windows.reduce((acc, curr) => acc + curr);
    }

    expire() {
        const now = this.now();
        let diff = now - this.time;
        if (diff < this.millisPerWindow) {
            // not expired yet
            return;
        }

        let idx = (this.generation + 1) % this.windows.length;

        // Handle zeroing out any windows that may have been "skipped"
        const skip = Math.min(Math.floor(diff / this.millisPerWindow) - 1, this.windows.length);
        if (skip > 0) {
            const r = Math.min(skip, this.windows.length - idx);
            this.windows.fill(0, idx, idx + r);
            this.windows.fill(0, 0, skip - r);
            idx = (idx + skip) % this.windows.length;
        }
        while (diff > this.millisPerWindow) {
            this.windows[idx] = 0;
            diff -= this.millisPerWindow;
            idx = (idx + 1) % this.windows.length;
        }

        this.time = now;
        this.generation = idx;
    }
}

test-budget.ts

import "jasmine";
import { strict as assert } from "assert";
import { RetryBudget } from "../src/retry";

describe("RetryBudget", () => {
    it("should be a leaky bucket", () => {
        const clock = new Clock();
        const budget = new RetryBudget({
            ttlInSeconds: 1,
            minPerSecond: 0,
            retryPercent: 1.0,
            now: () => clock.now(),
        });

        budget.deposit();

        // with 1s ttl, test the deposit is leaked
        clock.advance(2000);

        assert(!budget.withdraw());
    })

    it("should use the reserve", () => {
        const clock = new Clock();
        const budget = new RetryBudget({
            ttlInSeconds: 1,
            minPerSecond: 3,
            retryPercent: 1.0,
            now: () => clock.now(),
        });

        assert(budget.withdraw());
        assert(budget.withdraw());
        assert(budget.withdraw());

        assert(!budget.withdraw());
    })
});

class Clock {
    private time: number;
    constructor() {
        this.time = 0;
    }

    advance(millis: number) {
        this.time += millis;
    }

    now(): number {
        return this.time;
    }
}