Added redlock for locking resource (#40)

jainmohit2001 · Oct 11, 2023 · 2c5adcc · 2c5adcc
1 parent 5aced51
commit 2c5adcc
Show file tree

Hide file tree

Showing 5 changed files with 204 additions and 21 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -26,10 +26,12 @@
     "glob": "^10.3.9",
     "helmet": "^7.0.0",
     "ini": "^4.1.1",
+    "ioredis": "^4.28.5",
     "jsdom": "^22.1.0",
     "memcached": "^2.2.2",
     "raw-socket": "github:algj/node-raw-socket",
     "redis": "^4.6.10",
+    "redlock": "^4.2.0",
     "ts-node": "^10.9.1",
     "uuid": "^9.0.1",
     "winston": "^3.10.0"
@@ -46,6 +48,7 @@
     "@types/jsdom": "^21.1.2",
     "@types/memcached": "^2.2.7",
     "@types/node": "^20.4.2",
+    "@types/redlock": "^4.0.5",
     "@types/supertest": "^2.0.12",
     "@types/uuid": "^9.0.4",
     "@typescript-eslint/eslint-plugin": "^6.0.0",

diff --git a/src/27/algorithms/redis-sliding-window-counter.ts b/src/27/algorithms/redis-sliding-window-counter.ts
@@ -1,6 +1,7 @@
-import { createClient } from 'redis';
+import Client from 'ioredis';
 import { RateLimiter, RedisSlidingWindowCounterArgs } from '../types';
 import { Request, Response, NextFunction } from 'express';
+import Redlock from 'redlock';
 
 export type Counter = {
   /**
@@ -34,11 +35,15 @@ export type Counter = {
 
 export class RedisSlidingWindowCounterRateLimiter implements RateLimiter {
   threshold: number;
-  client: ReturnType<typeof createClient>;
 
-  constructor({ threshold, client }: RedisSlidingWindowCounterArgs) {
+  client: Client.Redis;
+
+  redlock: Redlock;
+
+  constructor({ threshold, client, redlock }: RedisSlidingWindowCounterArgs) {
     this.threshold = threshold;
     this.client = client;
+    this.redlock = redlock;
   }
 
   async handleRequest(req: Request, res: Response, next: NextFunction) {
@@ -52,23 +57,31 @@ export class RedisSlidingWindowCounterRateLimiter implements RateLimiter {
       return;
     }
 
-    ip = ip.replaceAll(':', '');
+    ip = ip.replaceAll(/:|\./g, '');
 
-    const value = (await this.client.json.get(ip)) as unknown;
-    const counter = value !== null ? (value as Counter) : null;
+    const lock = await this.redlock.acquire(ip + 'redlock', 5000);
+
+    const value = await this.client.get(ip);
+    const counter = value !== null ? (JSON.parse(value) as Counter) : null;
     const requestTimestamp = new Date().getTime();
 
     const currentWindow = Math.floor(new Date().getTime() / 1000) * 1000;
     const prevWindow = currentWindow - 1000;
 
     // If this is the first request from the given IP
     if (counter === null) {
-      await this.client.json.set(ip, '.', {
-        currentCounter: 1,
-        currentWindow: currentWindow,
-        prevCounter: 0,
-        prevWindow: prevWindow
-      });
+      await this.client.set(
+        ip,
+        JSON.stringify({
+          currentCounter: 1,
+          currentWindow: currentWindow,
+          prevCounter: 0,
+          prevWindow: prevWindow
+        })
+      );
+
+      await lock.unlock();
+
       next();
       return;
     }
@@ -97,14 +110,20 @@ export class RedisSlidingWindowCounterRateLimiter implements RateLimiter {
     // If the count is higher than the threshold, then reject the request
     if (count >= this.threshold) {
       // Update the counters for this IP
-      await this.client.json.set(ip, '.', counter);
+      await this.client.set(ip, JSON.stringify(counter));
+
+      await lock.unlock();
+
       res.status(429).send('Too many requests. Please try again later\n');
       return;
     }
 
     // update the current counter for this IP
     counter.currentCounter++;
-    await this.client.json.set(ip, '.', counter);
+    await this.client.set(ip, JSON.stringify(counter));
+
+    await lock.unlock();
+
     next();
   }
 

diff --git a/src/27/index.ts b/src/27/index.ts
@@ -9,7 +9,8 @@ import {
   SlidingWindowLogArgs,
   TokenBucketArgs
 } from './types';
-import { createClient } from 'redis';
+import Redlock from 'redlock';
+import Client from 'ioredis';
 
 program.addArgument(
   new Argument(
@@ -18,12 +19,14 @@ program.addArgument(
   ).choices(Object.values(RateLimiterType))
 );
 program.option('-p, --port <port>', 'Port on which server will start', '8080');
+program.option('--debug', 'Enable debugging');
+
 program.parse();
 
 const options = program.opts();
 const rateLimiterType = program.args[0] as RateLimiterType;
 const PORT = parseInt(options.port);
-const DEBUG = true;
+const DEBUG = options.debug;
 
 // Change the following default data to use different configuration.
 async function getRateLimiterArgs(
@@ -47,14 +50,25 @@ async function getRateLimiterArgs(
       return arg;
     }
     case RateLimiterType.REDIS_SLIDING_WINDOW_COUNTER: {
-      const client = createClient();
-      await client.connect();
+      const client = new Client();
       client.on('error', (err) => {
         if (DEBUG) {
           console.error(err);
         }
       });
-      const args: RedisSlidingWindowCounterArgs = { threshold: 1, client };
+
+      const redlock = new Redlock([client], {
+        driftFactor: 0.01,
+        retryCount: 10,
+        retryDelay: 200,
+        retryJitter: 200
+      });
+
+      const args: RedisSlidingWindowCounterArgs = {
+        threshold: 1,
+        client,
+        redlock
+      };
       return args;
     }
   }