Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion src/@types/rateLimit.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,19 @@ export interface TokenBucketOptions {
refillRate: number;
}

/**
* @type {number} windowSize - Size of each fixed window and the rolling window
* @type {number} capacity - Number of tokens a window can hold
*/
export interface SlidingWindowCounterOptions {
windowSize: number;
capacity: number;
}

// TODO: This will be a union type where we can specify Option types for other Rate Limiters
// Record<string, never> represents the empty object for alogorithms that don't require settings
// and might be able to be removed in the future.
export type RateLimiterOptions = TokenBucketOptions | Record<string, never>;
export type RateLimiterOptions =
| TokenBucketOptions
| SlidingWindowCounterOptions
| Record<string, never>;
9 changes: 7 additions & 2 deletions src/middleware/rateLimiterSetup.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import Redis from 'ioredis';
import { RateLimiterOptions, RateLimiterSelection } from '../@types/rateLimit';
import { RateLimiterOptions, RateLimiterSelection, TokenBucketOptions } from '../@types/rateLimit';
import SlidingWindowCounter from '../rateLimiters/slidingWindowCounter';
import TokenBucket from '../rateLimiters/tokenBucket';

/**
Expand All @@ -19,6 +20,8 @@ export default function setupRateLimiter(
switch (selection) {
case 'TOKEN_BUCKET':
// todo validate options
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
return new TokenBucket(options.bucketSize, options.refillRate, client);
break;
case 'LEAKY_BUCKET':
Expand All @@ -31,7 +34,9 @@ export default function setupRateLimiter(
throw new Error('Sliding Window Log has not be implemented.');
break;
case 'SLIDING_WINDOW_COUNTER':
throw new Error('Sliding Window Counter algonithm has not be implemented.');
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
return new SlidingWindowCounter(options.windowSize, options.capacity, client);
break;
default:
// typescript should never let us invoke this function with anything other than the options above
Expand Down
136 changes: 99 additions & 37 deletions src/rateLimiters/slidingWindowCounter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,12 @@ import { RateLimiter, RateLimiterResponse, RedisWindow } from '../@types/rateLim
* takeup in each.
*
* Whenever a user makes a request the following steps are performed:
* 1. Fixed minute windows are defined along with redis caches if previously undefined.
* 2. Rolling minute windows are defined or updated based on the timestamp of the new request.
* 1. Fixed windows are defined along with redis caches if previously undefined.
* 2. Rolling windows are defined or updated based on the timestamp of the new request.
* 3. Counter of the current fixed window is updated with the new request's token usage.
* 4. If a new minute interval is reached, the averaging formula is run to prevent fixed window's flaw
* of flooded requests around window borders
* (ex. 10 token capacity: 1m59s 10 reqs 2m2s 10 reqs)
* (ex. 1m windows, 10 token capacity: 1m59s 10 reqs 2m2s 10 reqs)
*/
class SlidingWindowCounter implements RateLimiter {
private windowSize: number;
Expand All @@ -24,10 +24,10 @@ class SlidingWindowCounter implements RateLimiter {
private client: Redis;

/**
* Create a new instance of a TokenBucket rate limiter that can be connected to any database store
* @param windowSize - size of each window in milliseconds (fixed and rolling)
* @param capacity - max capacity of tokens allowed per fixed window
* @param client - redis client where rate limiter will cache information
* Create a new instance of a SlidingWindowCounter rate limiter that can be connected to any database store
* @param windowSize size of each window in milliseconds (fixed and rolling)
* @param capacity max capacity of tokens allowed per fixed window
* @param client redis client where rate limiter will cache information
*/
constructor(windowSize: number, capacity: number, client: Redis) {
this.windowSize = windowSize;
Expand All @@ -38,12 +38,13 @@ class SlidingWindowCounter implements RateLimiter {
}

/**
* @function processRequest - current timestamp and number of tokens required for
* the request to go through are passed in. We first check if a window exists in the redis
* cache.
* @function processRequest - Sliding window counter algorithm to allow or block
* based on the depth/complexity (in amount of tokens) of incoming requests.
*
* If not, then fixedWindowStart is set as the current timestamp, and currentTokens
* is checked against capacity. If we have enough capacity for the request, we return
* First, checks if a window exists in the redis cache.
*
* If not, then `fixedWindowStart` is set as the current timestamp, and `currentTokens`
* is checked against `capacity`. If enough room exists for the request, returns
* success as true and tokens as how many tokens remain in the current fixed window.
*
* If a window does exist in the cache, we first check if the timestamp is greater than
Expand All @@ -66,6 +67,8 @@ class SlidingWindowCounter implements RateLimiter {
* @param {number} timestamp - time the request was recieved
* @param {number} [tokens=1] - complexity of the query for throttling requests
* @return {*} {Promise<RateLimiterResponse>}
* RateLimiterResponse: {success: boolean, tokens: number}
* (tokens represents the remaining available capacity of the window)
* @memberof SlidingWindowCounter
*/
async processRequest(
Expand All @@ -79,31 +82,90 @@ class SlidingWindowCounter implements RateLimiter {
// attempt to get the value for the uuid from the redis cache
const windowJSON = await this.client.get(uuid);

// // if the response is null, we need to create a window for the user
// if (windowJSON === null) {
// // rolling window is 1 minute long
// const rollingWindowEnd = timestamp + 60000;

// // grabs the actual minute from the timestamp to create fixed window
// const fixedWindowStart = timestamp - (timestamp % 10000);
// const fixedWindowEnd = fixedWindowStart + 60000;

// const newUserWindow: RedisWindow = {
// // conditionally set tokens depending on how many are requested compared to the capacity
// tokens: tokens > this.capacity ? this.capacity : this.capacity - tokens,
// timestamp,
// };

// // reject the request, not enough tokens could even be in the bucket
// if (tokens > this.capacity) {
// await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow));
// return { success: false, tokens: this.capacity };
// }
// await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow));
// return { success: true, tokens: newUserWindow.tokens };
// }

return { success: true, tokens: 0 };
// if the response is null, we need to create a window for the user
if (windowJSON === null) {
const newUserWindow: RedisWindow = {
// current and previous tokens represent how many tokens are in each window
currentTokens: tokens <= this.capacity ? tokens : 0,
previousTokens: 0,
fixedWindowStart: timestamp,
};

if (tokens <= this.capacity) {
await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow));
return { success: true, tokens: this.capacity - newUserWindow.currentTokens };
}

await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow));
// tokens property represents how much capacity remains
return { success: false, tokens: this.capacity };
}

// if the cache is populated

const window: RedisWindow = await JSON.parse(windowJSON);

const updatedUserWindow: RedisWindow = {
currentTokens: window.currentTokens,
previousTokens: window.previousTokens,
fixedWindowStart: window.fixedWindowStart,
};

// if request time is in a new window
if (window.fixedWindowStart && timestamp >= window.fixedWindowStart + this.windowSize) {
// if more than one window was skipped
if (timestamp >= window.fixedWindowStart + this.windowSize * 2) {
// if one or more windows was skipped, reset new window to be at current timestamp
updatedUserWindow.previousTokens = 0;
updatedUserWindow.currentTokens = 0;
updatedUserWindow.fixedWindowStart = timestamp;
} else {
updatedUserWindow.previousTokens = updatedUserWindow.currentTokens;
updatedUserWindow.currentTokens = 0;
updatedUserWindow.fixedWindowStart = window.fixedWindowStart + this.windowSize;
}
}

// assigned to avoid TS error, this var will never be used as 0
// var is declared here so that below can be inside a conditional for efficiency's sake
let rollingWindowProportion = 0;
let previousRollingTokens = 0;

if (updatedUserWindow.fixedWindowStart && updatedUserWindow.previousTokens) {
// proportion of rolling window present in previous window
rollingWindowProportion =
(this.windowSize - (timestamp - updatedUserWindow.fixedWindowStart)) /
this.windowSize;

// remove unecessary decimals, 0.xx is enough
// rollingWindowProportion -= rollingWindowProportion % 0.01;

// # of tokens present in rolling & previous window
previousRollingTokens = Math.floor(
updatedUserWindow.previousTokens * rollingWindowProportion
);
}

// # of tokens present in rolling and/or current window
// if previous tokens is null, previousRollingTokens will be 0
const rollingTokens = updatedUserWindow.currentTokens + previousRollingTokens;

// if request is allowed
if (tokens + rollingTokens <= this.capacity) {
updatedUserWindow.currentTokens += tokens;
await this.client.setex(uuid, keyExpiry, JSON.stringify(updatedUserWindow));
return {
success: true,
tokens: this.capacity - (updatedUserWindow.currentTokens + previousRollingTokens),
};
}

// if request is blocked
await this.client.setex(uuid, keyExpiry, JSON.stringify(updatedUserWindow));
return {
success: false,
tokens: this.capacity - (updatedUserWindow.currentTokens + previousRollingTokens),
};
}

/**
Expand Down
18 changes: 9 additions & 9 deletions test/rateLimiters/slidingWindowCounter.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ async function setTokenCountInClient(
await redisClient.set(uuid, JSON.stringify(value));
}

xdescribe('Test SlidingWindowCounter Rate Limiter', () => {
describe('Test SlidingWindowCounter Rate Limiter', () => {
beforeEach(async () => {
// init a mock redis cache
client = new RedisMock();
Expand Down Expand Up @@ -249,15 +249,15 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => {
const result = await limiter.processRequest(
user4,
timestamp + WINDOW_SIZE * 1.99,
4
10
);
expect(result.tokens).toBe(0);
expect(result.success).toBe(true);

// currentTokens (in current fixed window): 4
// previousTokens (in previous fixed window): 8
const count1 = await getWindowFromClient(client, user4);
expect(count1.currentTokens).toBe(4);
expect(count1.currentTokens).toBe(10);
expect(count1.previousTokens).toBe(8);
});
});
Expand Down Expand Up @@ -304,7 +304,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => {

// 3 + 8 * 1 = 11, above capacity (request should be blocked)
const result = await limiter.processRequest(user4, timestamp + WINDOW_SIZE, 3);
expect(result.tokens).toBe(10);
expect(result.tokens).toBe(2);
expect(result.success).toBe(false);

// currentTokens (in current fixed window): 0
Expand Down Expand Up @@ -332,7 +332,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => {
timestamp + WINDOW_SIZE * 1.25,
5
);
expect(result.tokens).toBe(10);
expect(result.tokens).toBe(4);
expect(result.success).toBe(false);

// currentTokens (in current fixed window): 0
Expand All @@ -358,7 +358,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => {

// 7 + 8 * .5 = 11, over capacity (request should be blocked)
const result = await limiter.processRequest(user4, timestamp + WINDOW_SIZE * 1.5, 7);
expect(result.tokens).toBe(10);
expect(result.tokens).toBe(6);
expect(result.success).toBe(false);

// currentTokens (in current fixed window): 0
Expand All @@ -383,7 +383,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => {

// 9 + 8 * .25 = 11, over capacity (request should be blocked)
const result = await limiter.processRequest(user4, timestamp + WINDOW_SIZE * 1.75, 9);
expect(result.tokens).toBe(10);
expect(result.tokens).toBe(8);
expect(result.success).toBe(false);

// currentTokens (in current fixed window): 0
Expand All @@ -407,7 +407,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => {

// 11 + 8 * .01 = 11, above capacity (request should be blocked)
const result = await limiter.processRequest(user4, timestamp + WINDOW_SIZE, 11);
expect(result.tokens).toBe(10);
expect(result.tokens).toBe(2);
expect(result.success).toBe(false);

// currentTokens (in current fixed window): 0
Expand Down Expand Up @@ -465,7 +465,7 @@ xdescribe('Test SlidingWindowCounter Rate Limiter', () => {
await (
await limiter.processRequest(user1, timestamp + WINDOW_SIZE, 4)
).tokens
).toBe(2);
).toBe(1);
// currentTokens (in current fixed window): 0
// previousTokens (in previous fixed window): 8
const count = await getWindowFromClient(client, user1);
Expand Down