fix running task count bug (opensearch-project#641)

Signed-off-by: Yaliang Wu <[email protected]> Signed-off-by: Yaliang Wu <[email protected]>
ylwu-amzn · Feb 17, 2023 · 0a146a2 · 0a146a2
1 parent 6f20220
commit 0a146a2
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/task/MLTaskManager.java b/plugin/src/main/java/org/opensearch/ml/task/MLTaskManager.java
@@ -81,18 +81,18 @@ public synchronized String checkLimitAndAddRunningTask(MLTask mlTask, Integer li
             runningTaskCount.set(0);
         }
         log.debug("Task id: {}, current running task {}: {}", mlTask.getTaskId(), mlTask.getTaskType(), runningTaskCount.get());
-        if (runningTaskCount.get() >= limit) {
-            String error = "exceed max running task limit";
-            log.info(error + " for task " + mlTask.getTaskId());
-            return error;
-        }
         if (contains(mlTask.getTaskId())) {
             getMLTask(mlTask.getTaskId()).setState(MLTaskState.RUNNING);
         } else {
+            if (runningTaskCount.get() >= limit) {
+                String error = "exceed max running task limit";
+                log.info(error + " for task " + mlTask.getTaskId());
+                return error;
+            }
             mlTask.setState(MLTaskState.RUNNING);
             add(mlTask);
+            runningTaskCount.incrementAndGet();
         }
-        runningTaskCount.incrementAndGet();
         return null;
     }
 
@@ -290,7 +290,7 @@ public void updateMLTask(
     ) {
         MLTaskCache taskCache = taskCaches.get(taskId);
         if (removeFromCache) {
-            taskCaches.remove(taskId);
+            remove(taskId);
         }
         if (taskCache == null) {
             listener.onFailure(new MLResourceNotFoundException("Can't find task"));
@@ -409,7 +409,7 @@ public void syncRunningLoadModelTasks(Map<String, Set<String>> runningLoadModelT
         if (staleTasks.size() > 0) {
             log.debug("remove stale load tasks : {}", Arrays.toString(staleTasks.toArray(new String[0])));
             for (String taskId : staleTasks) {
-                taskCaches.remove(taskId);
+                remove(taskId);
             }
         }
     }