[migration] fix bulk_indexing counting logic (#60)

[migration] fix bulk_indexing counting logic

Co-authored-by: Kassian Sun <kassiansun@outlook.com>

plugin/migration/model.go

@@ -38,7 +38,7 @@ type ElasticDataConfig struct {
 
 type ExecutionConfig struct {
 	TimeWindow []TimeWindowItem `json:"time_window"`
-	Nodes struct{
+	Nodes      struct {
 		Permit []ExecutionNode `json:"permit"`
 	} `json:"nodes"`
 }
@@ -78,7 +78,7 @@ type IndexInfo struct {
 	StoreSizeInBytes int    `json:"store_size_in_bytes"`
 }
 
-func (ii *IndexInfo) GetUniqueIndexName() string{
+func (ii *IndexInfo) GetUniqueIndexName() string {
 	return fmt.Sprintf("%s:%s", ii.Name, ii.DocType)
 }
 
@@ -93,13 +93,13 @@ type TaskCompleteState struct {
 	Error         string
 	ClearPipeline bool
 	PipelineIds   []string
-	SuccessDocs float64
-	ScrolledDocs float64
 	RunningPhase  int
-	TotalDocs interface{}
+	TotalDocs     int64
+	SuccessDocs   int64
+	ScrolledDocs  int64
 }
 
-type MajorTaskState struct{
+type MajorTaskState struct {
 	ScrolledDocs float64
 	IndexDocs    float64
 	Status       string

plugin/migration/pipeline.go

@@ -1262,46 +1262,64 @@ func (p *DispatcherProcessor) getTaskCompleteState(subTask *task2.Task) (*TaskCo
 	if cfg, ok = subTask.Config.(map[string]interface{}); !ok {
 		return nil, fmt.Errorf("got wrong config of task %v", *subTask)
 	}
-	totalDocs, err := util.MapStr(cfg).GetValue("source.doc_count")
+	totalDocsVal, err := util.MapStr(cfg).GetValue("source.doc_count")
 	if err != nil {
 		log.Errorf("failed to get source.doc_count, err: %v", err)
 		return nil, err
 	}
+	totalDocs, err := util.ExtractInt(totalDocsVal)
+	if err != nil {
+		log.Errorf("failed to extract source.doc_count, err: %v", err)
+		return nil, err
+	}
 
 	var (
-		indexDocs    float64
-		successDocs  float64
-		scrolledDocs interface{}
+		indexDocs    int64
+		successDocs  int64
+		scrolledDocs int64
 		state        TaskCompleteState
 	)
 	state.TotalDocs = totalDocs
 	state.PipelineIds = pids
+	var bulked, scrolled bool
 	for _, hit := range res.Hits.Hits {
+		if bulked && scrolled {
+			break
+		}
 		resultErr, _ := util.MapStr(hit.Source).GetValue("payload.pipeline.logging.result.error")
 		if errStr, ok := resultErr.(string); ok && errStr != "" {
 			state.Error = errStr
 			state.IsComplete = true
 			state.ClearPipeline = true
 		}
+		if !bulked {
 			for _, key := range []string{"payload.pipeline.logging.context.bulk_indexing.success.count", "payload.pipeline.logging.context.bulk_indexing.failure.count", "payload.pipeline.logging.context.bulk_indexing.invalid.count"} {
 				v, err := util.MapStr(hit.Source).GetValue(key)
 				if err == nil {
-				if fv, ok := v.(float64); ok {
+					bulked = true
+					if fv, err := util.ExtractInt(v); err == nil {
 						indexDocs += fv
 						if key == "payload.pipeline.logging.context.bulk_indexing.success.count" {
 							successDocs = fv
 							state.SuccessDocs = successDocs
 						}
-				}
 					} else {
-				break
+						log.Errorf("got %s but failed to extract, err: %v", key, err)
 					}
 				}
+			}
+		}
+
+		if !scrolled {
 			v, err := util.MapStr(hit.Source).GetValue("payload.pipeline.logging.context.es_scroll.scrolled_docs")
 			if err == nil {
-			scrolledDocs = v
-			if vv, ok := v.(float64); ok {
+				scrolled = true
+				if vv, err := util.ExtractInt(v); err == nil {
+					scrolledDocs = vv
 					state.ScrolledDocs = vv
+				} else {
+					log.Errorf("got payload.pipeline.logging.context.es_scroll.scrolled_docs but failed to extract, err: %v", err)
+				}
 			}
 		}
 	}