jhurliman / node-rate-limiter

A generic rate limiter for node.js. Useful for API clients, web crawling, or other tasks that need to be throttled
MIT License
1.51k stars 135 forks source link

Tiered Rate Limiting - Extending TokenBucket to support such functionality. #81

Closed lynniemagoo closed 1 year ago

lynniemagoo commented 3 years ago

I am calling an API that makes the specific requirements.

  1. You can only call said API 60 times / minute.
  2. You can only call said API 1000 times / day.

Obviously, you can satisfy the #1 by using new RateLimiter({"tokensPerInterval":60 ,"interval":"minute"}); You could then create a TokenBuffer for #2 with a parent as #1. However, the TokenBuffer does not add tokens up front and also 'drips' tokens in as time goes by.

Hence, I looked at another solution of using 2 tiered TokenBuckets with behavior similar to RateLimiter.

In order to make a TokenBucket behave initially like the RateLimiter, during construction, it is necessary to see the content at construction time.

This is currently done in RateLimiter as follows: this.tokenBucket.content = tokensPerInterval;

However, as this is a TokenBucket and not a RateLimiter, tokens are only added during calls to drip(). Ideally, to make TokenBucket behave in the same fashion as RateLimiter, a change to drip() was required.

Here's the code excerpt I modified such that you pass a value for milliseconds (necessary to allow extension of drip() for custom behavior by subclasses) and then use milliseconds accordingly.

this.drip(clock_1.getMilliseconds());

 /**
     * Add any new tokens to the bucket since the last drip.
     * @returns {Boolean} True if new tokens were added, otherwise false.
     */
    drip(milliseconds) {
        if (this.tokensPerInterval === 0) {
            const prevContent = this.content;
            this.content = this.bucketSize;
            return this.content > prevContent;
        }

        const now = milliseconds;
        const deltaMS = Math.max(now - this.lastDrip, 0);
        this.lastDrip = now;
        const dripAmount = deltaMS * (this.tokensPerInterval / this.interval);
        const prevContent = this.content;
        this.content = Math.min(this.content + dripAmount, this.bucketSize);
        return Math.floor(this.content) > Math.floor(prevContent);
    }

Next, you need to create the functionality you need with a custom drip() and initial seed of the content. Here's a class that does just that.

class RateLimiterEx extends TokenBucket {
    constructor({ bucketSize, tokensPerInterval, interval, parentBucket}) {
        super({bucketSize,  tokensPerInterval, interval, parentBucket});
        // initially fill the bucket.
        this.content = this.bucketSize;
    }

    // Overload drip() such that we add tokens only when the interval expires.
    drip(milliseconds) {
        if (this.tokensPerInterval === 0) {
            const prevContent = this.content;
            this.content = this.bucketSize;
            return this.content > prevContent;
        }
        // When timer expires, we simply reset content to the bucketSize.
        const now = milliseconds;
        const deltaMS = Math.max(now - this.lastDrip, 0);
        const prevContent = this.content;
        if (deltaMS >= this.interval) {
            this.content = this.bucketSize;
            this.lastDrip = now;
        }
        return Math.floor(this.content) > Math.floor(prevContent);
    }
}

With this implementation, fireImmediately is not supported. Also, there is a slight difference in the 'sleep' behavior such that TokenBucket will sleep 1000ms while awaiting a refresh (to ensure drip() is called every so often whereas RateLimiter sleeps for a larger amount of time.

Here's some sample code that I put together that illustrates the usage.

const { RateLimiter,TokenBucket } = require("limiter");

async function mainTB4() {
    optionsTier1 = {"tokensPerInterval":60, "interval":"minute"};
    optionsTier2 = {"tokensPerInterval":80, "interval":1000*80};
    const tb = new RateLimiterEx({...optionsTier1});
    const tb2 = new RateLimiterEx({...optionsTier2,"parentBucket":tb});
    let start = new Date().valueOf();
    let delta = start;
    for (let i=0;i<330;i++) {
        //await sleep(100);
        let now = new Date().valueOf();
        console.log("i:%d el:%d d:%d cr:%o", i,now - start, now - delta, await tb2.removeTokens(1)); 
        delta = now;
    }
}

mainTB4();

Attached please find sample code and changes to TokenBucket.js. RateLimiterEx.zip

I would welcome comments on this. IMO, this approach supports multiple levels of tiered limiting.

  1. You can only call said API 60 times / minute.
  2. You can only call said API 1000 times / month.
  3. You can only call said API 5000 times / year.