Merge pull request #80 from oslabs-beta/jd/slide

jondeweydev · web-flow · commit ea78c1354f42 · 2022-07-25T16:46:34.000-07:00
Sliding Window Counter Functionality Complete
diff --git a/src/@types/rateLimit.d.ts b/src/@types/rateLimit.d.ts
@@ -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>;
diff --git a/src/middleware/rateLimiterSetup.ts b/src/middleware/rateLimiterSetup.ts
@@ -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';
 
 /**
@@ -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':
@@ -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
diff --git a/src/rateLimiters/slidingWindowCounter.ts b/src/rateLimiters/slidingWindowCounter.ts
@@ -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;
@@ -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;
@@ -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
@@ -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(
@@ -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),
+        };
     }
 
     /**
diff --git a/test/rateLimiters/slidingWindowCounter.test.ts b/test/rateLimiters/slidingWindowCounter.test.ts
@@ -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();
@@ -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);
             });
         });
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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);
