From 5af5b2142e9aa522c42f81d8696931e0248ec70a Mon Sep 17 00:00:00 2001
From: Kassian Sun <kassiansun@outlook.com>
Date: Mon, 5 Jun 2023 17:09:12 +0800
Subject: [PATCH] [migration] fix task doc stats

---
 plugin/migration/comparison_api.go |  5 -----
 plugin/migration/migration_api.go  | 19 +++++++++++--------
 2 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/plugin/migration/comparison_api.go b/plugin/migration/comparison_api.go
index 231569ca..573f2804 100644
--- a/plugin/migration/comparison_api.go
+++ b/plugin/migration/comparison_api.go
@@ -142,11 +142,6 @@ func (h *APIHandler) getComparisonMajorTaskInfo(taskID string) (taskStats Cluste
 							},
 						},
 					},
-					{
-						"terms": util.MapStr{
-							"status": []string{task.StatusComplete, task.StatusError},
-						},
-					},
 				},
 			},
 		},
diff --git a/plugin/migration/migration_api.go b/plugin/migration/migration_api.go
index 146c98c6..56ce9778 100644
--- a/plugin/migration/migration_api.go
+++ b/plugin/migration/migration_api.go
@@ -89,23 +89,26 @@ func (h *APIHandler) getMigrationMajorTaskInfo(id string) (taskStats migration_m
 			continue
 		}
 
-		if subTask.Status == task.StatusRunning {
-			indexMigrationTaskIDs = append(indexMigrationTaskIDs, subTask.ID)
-			continue
-		}
-
 		cfg := migration_model.IndexMigrationTaskConfig{}
 		err = migration_util.GetTaskConfig(&subTask, &cfg)
 		if err != nil {
 			log.Errorf("failed to get task config, err: %v", err)
 			continue
 		}
-		indexDocs := migration_util.GetMapIntValue(taskLabels, "index_docs")
-		taskStats.IndexDocs += indexDocs
+
 		taskStats.SourceDocs += cfg.Source.DocCount
 		st := indexState[indexName]
-		st.IndexDocs += indexDocs
 		st.SourceDocs += cfg.Source.DocCount
+		indexState[indexName] = st
+
+		if subTask.Status == task.StatusRunning {
+			indexMigrationTaskIDs = append(indexMigrationTaskIDs, subTask.ID)
+			continue
+		}
+
+		indexDocs := migration_util.GetMapIntValue(taskLabels, "index_docs")
+		taskStats.IndexDocs += indexDocs
+		st.IndexDocs += indexDocs
 		if subTask.Status == task.StatusError {
 			st.ErrorPartitions += 1
 			taskStats.ErrorPartitions += 1