From c2911da09ea612fd456e14fdd85c342ac20fb138 Mon Sep 17 00:00:00 2001 From: Ruslan Gainutdinov Date: Wed, 17 Nov 2021 23:41:42 +0200 Subject: [PATCH 1/4] Limit multiple actions --- README.md | 12 +++++++----- src/index.test.ts | 16 ++++++++++++++++ src/index.ts | 33 +++++++++++++++++++-------------- 3 files changed, 42 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index b0de749..5cbdcbc 100644 --- a/README.md +++ b/README.md @@ -98,8 +98,10 @@ To run tests, you will need to have a Redis server running. You can do this by i ### Testing -- `yarn ci`: Runs the CI build, including linting, type checking, and tests. Requires [act](https://github.com/nektos/act) to run GitHub actions locally. -- `yarn lint`: Runs ESLint. -- `yarn test`: Runs Jest. -- `yarn typecheck`: Runs TypeScript, without emitting output. -- `yarn build`: Runs TypeScript and outputs to `./lib`. +Before running tests, make sure you have redis running. You can start redis in Docker, for example, by executing `docker run -it -p 6379:6379 -d --name=redis --restart=always redis:alpine` + +* `yarn ci`: Runs the CI build, including linting, type checking, and tests. Requires [act](https://github.com/nektos/act) to run GitHub actions locally. +* `yarn lint`: Runs ESLint. +* `yarn test`: Runs Jest. +* `yarn typecheck`: Runs TypeScript, without emitting output. +* `yarn build`: Runs TypeScript and outputs to `./lib`. diff --git a/src/index.test.ts b/src/index.test.ts index b649ab7..f7db43f 100644 --- a/src/index.test.ts +++ b/src/index.test.ts @@ -122,6 +122,22 @@ 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); + + // Should allow first action through. + setTime(1000); + expect(await limiter.wouldLimitWithInfo(id)).toEqual({ + actionsRemaining: 100, + blocked: false, + blockedDueToCount: false, + blockedDueToMinDifference: false, + millisecondsUntilAllowed: 0, + }); + expect(await limiter.limit(id, 2)).toBe(false); + }); + it('blocked actions count as actions', async () => { const options = { interval: 10, maxInInterval: 3 }; const limiter = await createLimiter(options); diff --git a/src/index.ts b/src/index.ts index 9391d74..cf49152 100644 --- a/src/index.ts +++ b/src/index.ts @@ -50,8 +50,8 @@ export class RateLimiter { * Attempts an action for the provided ID. Return information about whether the action was * allowed and why, and whether upcoming actions will be allowed. */ - async limitWithInfo(id: Id): Promise { - const timestamps = await this.getTimestamps(id, true); + async limitWithInfo(id: Id, count = 1): Promise { + const timestamps = await this.getTimestamps(id, count); return this.calculateInfo(timestamps); } @@ -60,15 +60,15 @@ export class RateLimiter { */ async wouldLimitWithInfo(id: Id): Promise { 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. */ - async limit(id: Id): Promise { - return (await this.limitWithInfo(id)).blocked; + async limit(id: Id, count = 1): Promise { + 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> { 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> { 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,7 +236,7 @@ export class RedisRateLimiter extends RateLimiter { protected async getTimestamps( id: Id, - addNewTimestamp: boolean, + addNewTimestamps: number, ): Promise> { const now = getCurrentMicroseconds(); const key = this.makeKey(id); @@ -241,8 +244,8 @@ export class RedisRateLimiter extends RateLimiter { 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,9 @@ export class RedisRateLimiter extends RateLimiter { batch.exec((err, result) => { if (err) return reject(err); - const zRangeOutput = (addNewTimestamp ? result[2] : result[1]) as Array; + const zRangeOutput = (addNewTimestamps > 0 + ? result[2] + : result[1]) as Array; const zRangeResult = this.getZRangeResult(zRangeOutput); const timestamps = this.extractTimestampsFromZRangeResult(zRangeResult); return resolve(timestamps); From 6abe54cf16df47e7d14466ee2dea24da9c562037 Mon Sep 17 00:00:00 2001 From: Ruslan Gainutdinov Date: Wed, 17 Nov 2021 23:48:26 +0200 Subject: [PATCH 2/4] Improve test2 --- src/index.test.ts | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/src/index.test.ts b/src/index.test.ts index f7db43f..7594dba 100644 --- a/src/index.test.ts +++ b/src/index.test.ts @@ -126,16 +126,33 @@ describe('RateLimiter implementations', () => { const options = { interval: 10, maxInInterval: 100 }; const limiter = await createLimiter(options); - // Should allow first action through. setTime(1000); + + // All actions available for consuming expect(await limiter.wouldLimitWithInfo(id)).toEqual({ - actionsRemaining: 100, + actionsRemaining: options.maxInInterval, blocked: false, blockedDueToCount: false, blockedDueToMinDifference: false, millisecondsUntilAllowed: 0, }); - expect(await limiter.limit(id, 2)).toBe(false); + + // 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 () => { From 678cdd925a9f5a281d8997bf43fee97b0d7de882 Mon Sep 17 00:00:00 2001 From: Ruslan Gainutdinov Date: Thu, 18 Nov 2021 22:22:49 +0200 Subject: [PATCH 3/4] Fix docker instruction in readme --- README.md | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 5cbdcbc..3bddb48 100644 --- a/README.md +++ b/README.md @@ -94,14 +94,16 @@ 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: -### Testing +- 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. -Before running tests, make sure you have redis running. You can start redis in Docker, for example, by executing `docker run -it -p 6379:6379 -d --name=redis --restart=always redis:alpine` +### Testing -* `yarn ci`: Runs the CI build, including linting, type checking, and tests. Requires [act](https://github.com/nektos/act) to run GitHub actions locally. -* `yarn lint`: Runs ESLint. -* `yarn test`: Runs Jest. -* `yarn typecheck`: Runs TypeScript, without emitting output. -* `yarn build`: Runs TypeScript and outputs to `./lib`. +- `yarn ci`: Runs the CI build, including linting, type checking, and tests. Requires [act](https://github.com/nektos/act) to run GitHub actions locally. +- `yarn lint`: Runs ESLint. +- `yarn test`: Runs Jest. +- `yarn typecheck`: Runs TypeScript, without emitting output. +- `yarn build`: Runs TypeScript and outputs to `./lib`. From 471bcc47eba00b231ffbea162929e3b7bf186f75 Mon Sep 17 00:00:00 2001 From: Ruslan Gainutdinov Date: Thu, 18 Nov 2021 22:48:18 +0200 Subject: [PATCH 4/4] Fix redis batch result --- README.md | 4 ++-- src/index.test.ts | 8 ++++++++ src/index.ts | 8 +++----- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 3bddb48..6eda36f 100644 --- a/README.md +++ b/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` - 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` - 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` - 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` - 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` - 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` - Returns info about what would happened if an action were attempted and why. Does not "count" as an action. `RateLimitInfo` contains the following properties: diff --git a/src/index.test.ts b/src/index.test.ts index 7594dba..8adeb66 100644 --- a/src/index.test.ts +++ b/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. diff --git a/src/index.ts b/src/index.ts index cf49152..6d608c0 100644 --- a/src/index.ts +++ b/src/index.ts @@ -47,7 +47,7 @@ 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, count = 1): Promise { @@ -65,7 +65,7 @@ export class RateLimiter { } /** - * 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, count = 1): Promise { return (await this.limitWithInfo(id, count)).blocked; @@ -254,9 +254,7 @@ export class RedisRateLimiter extends RateLimiter { batch.exec((err, result) => { if (err) return reject(err); - const zRangeOutput = (addNewTimestamps > 0 - ? result[2] - : result[1]) as Array; + const zRangeOutput = result[1 + addNewTimestamps] as Array; const zRangeResult = this.getZRangeResult(zRangeOutput); const timestamps = this.extractTimestampsFromZRangeResult(zRangeResult); return resolve(timestamps);