From 0a146a2cdcd77de6a9653da3ccce749017e2b8ce Mon Sep 17 00:00:00 2001
From: Yaliang Wu <ylwu@amazon.com>
Date: Wed, 21 Dec 2022 00:13:03 +0000
Subject: [PATCH] fix running task count bug (#641)

Signed-off-by: Yaliang Wu <ylwu@amazon.com>

Signed-off-by: Yaliang Wu <ylwu@amazon.com>
---
 .../org/opensearch/ml/task/MLTaskManager.java    | 16 ++++++++--------
 1 file changed, 8 insertions(+), 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
index cd1acd4507..b689555a50 100644
--- 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);
             }
         }
     }