Limit multiple actions

README.md

@@ -64,9 +64,9 @@ app.use(function (req, res, next) {
 
 All methods take an `Id`, which should be of type `number | string`. Commonly, this will be a user's id.
 
-- `limit(id: Id): Promise<boolean>` - Attempt to perform an action. Returns `false` if the action should be allowed, and `true` if the action should be blocked.
+- `limit(id: Id, count = 1): Promise<boolean>` - Attempts one or more actions for the provided ID. Returns `false` if the action(s) should be allowed, and `true` if the action(s) should be blocked.
 - `wouldLimit(id: Id): Promise<boolean>` - Return what would happen if an action were attempted. Returns `false` if an action would not have been blocked, and `true` if an action would have been blocked. Does not "count" as an action.
-- `limitWithInfo(id: Id): Promise<RateLimitInfo>` - Attempt to perform an action. Returns whether the action should be blocked, as well as additional information about why it was blocked and how long the user must wait.
+- `limitWithInfo(id: Id, count = 1): Promise<RateLimitInfo>` - Attempts one or more actions for the provided ID. Return information about whether the action(s) were allowed and why, and whether upcoming actions will be allowed.
 - `wouldLimitWithInfo(id: Id): Promise<RateLimitInfo>` - Returns info about what would happened if an action were attempted and why. Does not "count" as an action.
 
 `RateLimitInfo` contains the following properties:
@@ -94,7 +94,11 @@ All methods take an `Id`, which should be of type `number | string`. Commonly, t
 
 Install dependencies with `yarn`.
 
-To run tests, you will need to have a Redis server running. You can do this by installing Redis, and running `redis-server`. Alternatively, you can run the CI build, which includes tests, by installing [act](https://github.com/nektos/act). This requires Docker to be running - on MacOS that means running `Docker.app` from your `Applications` folder.
+To run tests, you will need to have a Redis server running. There are a few methods to do this:
+
+- Install Redis (using something like homebrew), and then run `redis-server`.
+- Run Redis from Docker: `docker run -it -p 6379:6379 -d --name=redis --restart=always redis:alpine`
+- Run the CI build, which includes tests, by installing [act](https://github.com/nektos/act). This requires Docker to be running - on MacOS that means running `Docker.app` from your `Applications` folder.
 
 ### Testing
 

src/index.test.ts

@@ -104,6 +104,14 @@ describe('RateLimiter implementations', () => {
 
       // Should allow first action through.
       setTime(0);
+      // Should have 2 actions available at start
+      expect(await limiter.wouldLimitWithInfo(id)).toEqual({
+        actionsRemaining: 2,
+        blocked: false,
+        blockedDueToCount: false,
+        blockedDueToMinDifference: false,
+        millisecondsUntilAllowed: 0,
+      });
       expect(await limiter.limit(id)).toBe(false);
 
       // Should allow second action through.
@@ -122,6 +130,39 @@ describe('RateLimiter implementations', () => {
       expect(await limiter.limit(id)).toBe(false);
     });
 
+    it('can lock multiple actions at once', async () => {
+      const options = { interval: 10, maxInInterval: 100 };
+      const limiter = await createLimiter(options);
+
+      setTime(1000);
+
+      // All actions available for consuming
+      expect(await limiter.wouldLimitWithInfo(id)).toEqual({
+        actionsRemaining: options.maxInInterval,
+        blocked: false,
+        blockedDueToCount: false,
+        blockedDueToMinDifference: false,
+        millisecondsUntilAllowed: 0,
+      });
+
+      // Lock all at once
+      expect(await limiter.limit(id, options.maxInInterval)).toBe(false);
+      // Nothing can be locked
+      expect(await limiter.limit(id)).toBe(true);
+
+      // Wait for window to pass
+      setTime(1000 + options.interval);
+      // All available again
+      expect(await limiter.limit(id, options.maxInInterval)).toBe(false);
+      expect(await limiter.limit(id)).toBe(true);
+
+      // Wait for half a window more
+      setTime(1000 + options.interval + options.interval / 2);
+      // Umm, should it have half of maxInInterval?
+      expect(await limiter.limit(id, options.maxInInterval / 2)).toBe(false);
+      expect(await limiter.limit(id)).toBe(true);
+    });
+
     it('blocked actions count as actions', async () => {
       const options = { interval: 10, maxInInterval: 3 };
       const limiter = await createLimiter(options);

src/index.ts

@@ -47,11 +47,11 @@ export class RateLimiter {
   }
 
   /**
-   * Attempts an action for the provided ID. Return information about whether the action was
+   * Attempts one or more actions for the provided ID. Return information about whether the action(s) were
    * allowed and why, and whether upcoming actions will be allowed.
    */
-  async limitWithInfo(id: Id): Promise<RateLimitInfo> {
-    const timestamps = await this.getTimestamps(id, true);
+  async limitWithInfo(id: Id, count = 1): Promise<RateLimitInfo> {
+    const timestamps = await this.getTimestamps(id, count);
     return this.calculateInfo(timestamps);
   }
 
@@ -60,15 +60,15 @@ export class RateLimiter {
    */
   async wouldLimitWithInfo(id: Id): Promise<RateLimitInfo> {
     const currentTimestamp = getCurrentMicroseconds();
-    const existingTimestamps = await this.getTimestamps(id, false);
+    const existingTimestamps = await this.getTimestamps(id, 0);
     return this.calculateInfo([...existingTimestamps, currentTimestamp]);
   }
 
   /**
-   * Attempts an action for the provided ID. Returns whether it was blocked.
+   * Attempts one or more actions for the provided ID. Return whether any actions were blocked.
    */
-  async limit(id: Id): Promise<boolean> {
-    return (await this.limitWithInfo(id)).blocked;
+  async limit(id: Id, count = 1): Promise<boolean> {
+    return (await this.limitWithInfo(id, count)).blocked;
   }
 
   /**
@@ -87,11 +87,11 @@ export class RateLimiter {
 
   /**
    * Returns the list of timestamps of actions attempted within `interval` for the provided ID. If
-   * `addNewTimestamp` flag is set, adds a new action with the current microsecond timestamp.
+   * `addNewTimestamps` is set, adds a new actions with the current microsecond timestamp and more, incrementing by 1.
    */
   protected async getTimestamps(
     _id: Id,
-    _addNewTimestamp: boolean,
+    _addNewTimestamps: number,
   ): Promise<Array<Microseconds>> {
     return Promise.reject(new Error('Not implemented'));
   }
@@ -159,14 +159,17 @@ export class InMemoryRateLimiter extends RateLimiter {
     }
   }
 
-  protected async getTimestamps(id: Id, addNewTimestamp: boolean) {
+  protected async getTimestamps(
+    id: Id,
+    addNewTimestamps: number,
+  ): Promise<Array<Microseconds>> {
     const currentTimestamp = getCurrentMicroseconds();
     // Update the stored timestamps, including filtering out old ones, and adding the new one.
     const clearBefore = currentTimestamp - this.interval;
     const storedTimestamps = (this.storage[id] || []).filter((t) => t > clearBefore);
 
-    if (addNewTimestamp) {
-      storedTimestamps.push(currentTimestamp);
+    for (let i = 0; i < addNewTimestamps; i++) {
+      storedTimestamps.push(currentTimestamp + i);
 
       // Set a new TTL, and cancel the old one, if present.
       const ttl = this.ttls[id];
@@ -233,16 +236,16 @@ export class RedisRateLimiter extends RateLimiter {
 
   protected async getTimestamps(
     id: Id,
-    addNewTimestamp: boolean,
+    addNewTimestamps: number,
   ): Promise<Array<Microseconds>> {
     const now = getCurrentMicroseconds();
     const key = this.makeKey(id);
     const clearBefore = now - this.interval;
 
     const batch = this.client.multi();
     batch.zremrangebyscore(key, 0, clearBefore);
-    if (addNewTimestamp) {
-      batch.zadd(key, String(now), uuid());
+    for (let i = 0; i < addNewTimestamps; i++) {
+      batch.zadd(key, String(now + i), uuid());
     }
     batch.zrange(key, 0, -1, 'WITHSCORES');
     batch.expire(key, this.ttl);
@@ -251,7 +254,7 @@ export class RedisRateLimiter extends RateLimiter {
       batch.exec((err, result) => {
         if (err) return reject(err);
 
-        const zRangeOutput = (addNewTimestamp ? result[2] : result[1]) as Array<unknown>;
+        const zRangeOutput = result[1 + addNewTimestamps] as Array<unknown>;
         const zRangeResult = this.getZRangeResult(zRangeOutput);
         const timestamps = this.extractTimestampsFromZRangeResult(zRangeResult);
         return resolve(timestamps);