PharmaLedger-IMI/epi-workspace#1042 added readiness probe

OpenDSU · Nov 16, 2022 · d29f848 · d29f848
1 parent 31c3177
commit d29f848
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 29 deletions.
diff --git a/index.js b/index.js
@@ -7,8 +7,6 @@ process.on('uncaughtException', err => {
 const httpWrapper = require('./libs/http-wrapper');
 const Server = httpWrapper.Server;
 
-const TokenBucket = require('./libs/TokenBucket');
-const START_TOKENS = 6000000;
 const CHECK_FOR_RESTART_COMMAND_FILE_INTERVAL = 500;
 
 (function loadDefaultComponents(){
@@ -41,8 +39,6 @@ function HttpServer({ listeningPort, rootFolder, sslConfig, dynamicPort, restart
 	}
 
 	let port = listeningPort || 8080;
-	const tokenBucket = new TokenBucket(START_TOKENS, 1, 10);
-
 	const conf =  require('./config').getConfig();
 	const server = new Server(sslConfig);
 	server.config = conf;
@@ -127,31 +123,6 @@ function HttpServer({ listeningPort, rootFolder, sslConfig, dynamicPort, restart
 			next();
 		});
 
-		if (conf.preventRateLimit !== true) {
-			server.use(function (req, res, next) {
-				const ip = res.socket.remoteAddress;
-				tokenBucket.takeToken(ip, tokenBucket.COST_MEDIUM, function (err, remainedTokens) {
-					res.setHeader('X-RateLimit-Limit', tokenBucket.getLimitByCost(tokenBucket.COST_MEDIUM));
-					res.setHeader('X-RateLimit-Remaining', tokenBucket.getRemainingTokenByCost(remainedTokens, tokenBucket.COST_MEDIUM));
-
-					if (err) {
-						if (err === TokenBucket.ERROR_LIMIT_EXCEEDED) {
-							res.statusCode = 429;
-						} else {
-							res.statusCode = 500;
-						}
-
-						res.end();
-						return;
-					}
-
-					next();
-				});
-			});
-		} else {
-			logger.trace(`Rate limit mechanism disabled!`);
-		}
-
 		server.options('/*', function (req, res) {
 			const headers = {};
 			// IE8 does not allow domains to be specified, just the *
@@ -176,6 +147,7 @@ function HttpServer({ listeningPort, rootFolder, sslConfig, dynamicPort, restart
         function addRootMiddlewares() {
 			const LoggerMiddleware = require('./middlewares/logger');
 			const AuthorisationMiddleware = require('./middlewares/authorisation');
+			const Throttler = require('./middlewares/throttler');
 			const OAuth = require('./middlewares/oauth');
 			const ResponseHeaderMiddleware = require('./middlewares/responseHeader');
 			const genericErrorMiddleware = require('./middlewares/genericErrorMiddleware');
@@ -187,6 +159,7 @@ function HttpServer({ listeningPort, rootFolder, sslConfig, dynamicPort, restart
 
 			genericErrorMiddleware(server);
 			requestEnhancements(server);
+			Throttler(server);
 
             if(conf.enableJWTAuthorisation) {
                 new AuthorisationMiddleware(server);

diff --git a/middlewares/throttler/index.js b/middlewares/throttler/index.js
@@ -0,0 +1,55 @@
+const TokenBucket = require("../../libs/TokenBucket");
+
+function Throttler(server){
+	const logger = $$.getLogger("apihub", "Throttler");
+	const START_TOKENS = 6000000;
+	const tokenBucket = new TokenBucket(START_TOKENS, 1, 10);
+	let remainingTokens = START_TOKENS;
+	const conf =  require('../../config').getConfig();
+	 function throttlerHandler(req, res, next) {
+		 const ip = res.socket.remoteAddress;
+		 tokenBucket.takeToken(ip, tokenBucket.COST_MEDIUM, function (err, remainedTokens) {
+			 res.setHeader('X-RateLimit-Limit', tokenBucket.getLimitByCost(tokenBucket.COST_MEDIUM));
+			 res.setHeader('X-RateLimit-Remaining', tokenBucket.getRemainingTokenByCost(remainedTokens, tokenBucket.COST_MEDIUM));
+			 remainingTokens = remainedTokens;
+			 if (err) {
+				 if (err === TokenBucket.ERROR_LIMIT_EXCEEDED) {
+					 res.statusCode = 429;
+				 } else {
+					 res.statusCode = 500;
+				 }
+
+				 res.end();
+				 return;
+			 }
+
+			 next();
+		 });
+	 }
+
+	 function readyProbeHandler(req, res) {
+		 const stats = {
+			 remainingTokens: tokenBucket.getRemainingTokenByCost(remainingTokens, tokenBucket.COST_MEDIUM),
+			 tokenLimit: tokenBucket.getLimitByCost(tokenBucket.COST_MEDIUM)
+		 }
+
+		 res.setHeader('Content-Type', 'application/json');
+		 res.statusCode = 200;
+		 res.write(JSON.stringify(stats));
+		 res.end();
+	 }
+
+	if (conf.preventRateLimit !== true) {
+		server.use(throttlerHandler);
+		server.get("/ready-probe", readyProbeHandler);
+	} else {
+		logger.trace(`Rate limit mechanism disabled!`);
+		server.get("/ready-probe", function (req, res) {
+			res.statusCode = 200;
+			res.write("Server ready");
+			res.end();
+		});
+	}
+}
+
+module.exports = Throttler;