Add decreaseCount to concurrency monitor

Author: pierotofy

libs/concurrencyMonitor.js

@@ -0,0 +1,48 @@
+/**
+ *  ClusterODM - A reverse proxy, load balancer and task tracker for NodeODM
+ *  Copyright (C) 2018-present MasseranoLabs LLC
+ *
+ *  This program is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU Affero General Public License as
+ *  published by the Free Software Foundation, either version 3 of the
+ *  License, or (at your option) any later version.
+ *
+ *  This program is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU Affero General Public License for more details.
+ *
+ *  You should have received a copy of the GNU Affero General Public License
+ *  along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+let userTasks = {};
+const LIMIT_WINDOW = 60 * 1000; // 1 minute
+
+module.exports = {
+    checkCommitLimitReached: function(maxConcurrentTasks, userToken){
+        if (maxConcurrentTasks === 0) return true;
+        if (!maxConcurrentTasks) return false;
+
+        if (!userToken) userToken = "default";
+        let now = new Date().getTime();
+
+        // Remove expired items
+        for (let t in userTasks){
+            userTasks[t] = userTasks[t].filter(tm => now - tm <= LIMIT_WINDOW);
+            if (userTasks[t].length === 0) delete(userTasks[t]);
+        }
+
+        if (userTasks[userToken] === undefined) userTasks[userToken] = [];
+
+        userTasks[userToken].push(now);
+        return userTasks[userToken].length > maxConcurrentTasks;
+    },
+
+    decreaseCount: function(userToken){
+        if (!userToken) userToken = "default";
+        if (userTasks[userToken] !== undefined && userTasks[userToken].length > 0){
+            userTasks[userToken].shift();
+        }
+    },
+};

libs/proxy.js

@@ -36,6 +36,7 @@ const async = require('async');
 const odmOptions = require('./odmOptions');
 const asrProvider = require('./asrProvider');
 const floodMonitor = require('./floodMonitor');
+const concurrencyMonitor = require('./concurrencyMonitor');
 const AWS = require('aws-sdk');
 
 module.exports = {
@@ -209,13 +210,15 @@ module.exports = {
                         const taskId = taskInfo.uuid;
 
                         asrProvider.onCommit(taskId, 10 * 1000);
-
+                        
                         // Add reference to S3 path if necessary
                         if (asrProvider.downloadsPath()){
                             taskInfo.s3Path = asrProvider.downloadsPath();
                         }
-
+                        
                         const token = await routetable.lookupToken(taskId);
+                        concurrencyMonitor.decreaseCount(token);
+                        
                         try{
                             cloudProvider.taskFinished(token, taskInfo);
                         }catch(e){
@@ -370,11 +373,17 @@ module.exports = {
                             asrProvider.cleanup(taskId);
                         };
 
+                        if (concurrencyMonitor.checkCommitLimitReached(limits.maxConcurrentTasks, query.token)){
+                            die(`Reached maximum number of concurrent tasks, please wait until other tasks have finished, then restart the task.`);
+                            return;
+                        }
+
                         if (await maxConcurrencyLimitReached(limits.maxConcurrentTasks, query.token)){
-                            die(`Reached maximum number of concurrent tasks: ${limits.maxConcurrentTasks}. Please wait until other tasks have finished, then restart the task.`);
+                            die(`Reached maximum number of concurrent tasks. Please wait until other tasks have finished, then restart the task.`);
                             return;
                         }
 
+
                         floodMonitor.recordTaskCommit(query.token);
                         utils.markTaskAsCommitted(taskId);
 
@@ -456,6 +465,8 @@ module.exports = {
                     });
                     busboy.on('finish', async function() {
                         if (taskId){
+                            concurrencyMonitor.decreaseCount(query.token);
+
                             let node = await routetable.lookupNode(taskId);
                             if (node){
                                 overrideRequest(req, node, query, pathname);