From da56880d0f763ef757ec53723203ec9a5f76ac76 Mon Sep 17 00:00:00 2001
From: wayblink <anyang.wang@zilliz.com>
Date: Wed, 3 Jul 2024 19:08:09 +0800
Subject: [PATCH] fix: Avoid datarace in clustering compaction (#34288)

#34289

Signed-off-by: wayblink <anyang.wang@zilliz.com>
---
 internal/datacoord/compaction_task.go           | 12 ++++++++++++
 .../datacoord/compaction_task_clustering.go     | 17 +++++++----------
 internal/proto/data_coord.proto                 |  1 +
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/internal/datacoord/compaction_task.go b/internal/datacoord/compaction_task.go
index 57e52f21c3..6cfdcb9af8 100644
--- a/internal/datacoord/compaction_task.go
+++ b/internal/datacoord/compaction_task.go
@@ -100,3 +100,15 @@ func setStartTime(startTime int64) compactionTaskOpt {
 		task.StartTime = startTime
 	}
 }
+
+func setRetryTimes(retryTimes int32) compactionTaskOpt {
+	return func(task *datapb.CompactionTask) {
+		task.RetryTimes = retryTimes
+	}
+}
+
+func setLastStateStartTime(lastStateStartTime int64) compactionTaskOpt {
+	return func(task *datapb.CompactionTask) {
+		task.LastStateStartTime = lastStateStartTime
+	}
+}
diff --git a/internal/datacoord/compaction_task_clustering.go b/internal/datacoord/compaction_task_clustering.go
index 4dc446a77c..4f7c9b5cc5 100644
--- a/internal/datacoord/compaction_task_clustering.go
+++ b/internal/datacoord/compaction_task_clustering.go
@@ -47,10 +47,9 @@ const (
 
 type clusteringCompactionTask struct {
 	*datapb.CompactionTask
-	plan                *datapb.CompactionPlan
-	result              *datapb.CompactionPlanResult
-	span                trace.Span
-	lastUpdateStateTime int64
+	plan   *datapb.CompactionPlan
+	result *datapb.CompactionPlanResult
+	span   trace.Span
 
 	meta             CompactionMeta
 	sessions         SessionManager
@@ -66,24 +65,22 @@ func (t *clusteringCompactionTask) Process() bool {
 		log.Warn("fail in process task", zap.Error(err))
 		if merr.IsRetryableErr(err) && t.RetryTimes < taskMaxRetryTimes {
 			// retry in next Process
-			t.RetryTimes = t.RetryTimes + 1
+			t.updateAndSaveTaskMeta(setRetryTimes(t.RetryTimes + 1))
 		} else {
 			log.Error("task fail with unretryable reason or meet max retry times", zap.Error(err))
-			t.State = datapb.CompactionTaskState_failed
-			t.FailReason = err.Error()
+			t.updateAndSaveTaskMeta(setState(datapb.CompactionTaskState_failed), setFailReason(err.Error()))
 		}
 	}
 	// task state update, refresh retry times count
 	currentState := t.State.String()
 	if currentState != lastState {
-		t.RetryTimes = 0
 		ts := time.Now().UnixMilli()
-		lastStateDuration := ts - t.lastUpdateStateTime
+		t.updateAndSaveTaskMeta(setRetryTimes(0), setLastStateStartTime(ts))
+		lastStateDuration := ts - t.GetLastStateStartTime()
 		log.Info("clustering compaction task state changed", zap.String("lastState", lastState), zap.String("currentState", currentState), zap.Int64("elapse", lastStateDuration))
 		metrics.DataCoordCompactionLatency.
 			WithLabelValues(fmt.Sprint(typeutil.IsVectorType(t.GetClusteringKeyField().DataType)), datapb.CompactionType_ClusteringCompaction.String(), lastState).
 			Observe(float64(lastStateDuration))
-		t.lastUpdateStateTime = ts
 
 		if t.State == datapb.CompactionTaskState_completed {
 			t.updateAndSaveTaskMeta(setEndTime(ts))
diff --git a/internal/proto/data_coord.proto b/internal/proto/data_coord.proto
index 81af5196a6..c2b5a8e5e2 100644
--- a/internal/proto/data_coord.proto
+++ b/internal/proto/data_coord.proto
@@ -913,6 +913,7 @@ message CompactionTask{
   int64 prefer_segment_rows = 22;
   int64 analyzeTaskID = 23;
   int64 analyzeVersion = 24;
+  int64 lastStateStartTime = 25;
 }
 
 message PartitionStatsInfo {