Fix clustering compaction task leak

Signed-off-by: Cai Zhang <[email protected]>
xiaocai2333 · Oct 12, 2024 · f158228 · f158228
1 parent 1d9c746
commit f158228
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 0 deletions.
diff --git a/internal/datacoord/compaction_task_clustering.go b/internal/datacoord/compaction_task_clustering.go
@@ -273,6 +273,11 @@ func (t *clusteringCompactionTask) processExecuting() error {
 }
 
 func (t *clusteringCompactionTask) processMetaSaved() error {
+	if err := t.sessions.DropCompactionPlan(t.GetNodeID(), &datapb.DropCompactionPlanRequest{
+		PlanID: t.GetPlanID(),
+	}); err != nil {
+		log.Warn("clusteringCompactionTask processFailedOrTimeout unable to drop compaction plan", zap.Int64("planID", t.GetPlanID()), zap.Error(err))
+	}
 	return t.updateAndSaveTaskMeta(setState(datapb.CompactionTaskState_indexing))
 }
 
@@ -363,6 +368,13 @@ func (t *clusteringCompactionTask) resetSegmentCompacting() {
 
 func (t *clusteringCompactionTask) processFailedOrTimeout() error {
 	log.Info("clean task", zap.Int64("triggerID", t.GetTriggerID()), zap.Int64("planID", t.GetPlanID()), zap.String("state", t.GetState().String()))
+
+	if err := t.sessions.DropCompactionPlan(t.GetNodeID(), &datapb.DropCompactionPlanRequest{
+		PlanID: t.GetPlanID(),
+	}); err != nil {
+		log.Warn("clusteringCompactionTask processFailedOrTimeout unable to drop compaction plan", zap.Int64("planID", t.GetPlanID()), zap.Error(err))
+	}
+
 	// revert segments meta
 	var operators []UpdateOperator
 	// revert level of input segments

diff --git a/internal/datacoord/compaction_task_clustering_test.go b/internal/datacoord/compaction_task_clustering_test.go
@@ -105,6 +105,7 @@ func (s *ClusteringCompactionTaskSuite) TestClusteringCompactionSegmentMetaChang
 		},
 	})
 	s.session.EXPECT().Compaction(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	s.session.EXPECT().DropCompactionPlan(mock.Anything, mock.Anything).Return(nil)
 
 	task := s.generateBasicTask(false)
 
@@ -370,6 +371,7 @@ func (s *ClusteringCompactionTaskSuite) TestProcessExecuting() {
 				},
 			},
 		}, nil).Once()
+		s.session.EXPECT().DropCompactionPlan(mock.Anything, mock.Anything).Return(nil).Once()
 		s.Equal(false, task.Process())
 		s.Equal(datapb.CompactionTaskState_indexing, task.GetState())
 	})
@@ -403,6 +405,8 @@ func (s *ClusteringCompactionTaskSuite) TestProcessExecuting() {
 				},
 			},
 		}, nil).Once()
+		// DropCompactionPlan fail
+		s.session.EXPECT().DropCompactionPlan(mock.Anything, mock.Anything).Return(merr.WrapErrNodeNotFound(1)).Once()
 		s.Equal(false, task.Process())
 		s.Equal(datapb.CompactionTaskState_indexing, task.GetState())
 	})
@@ -438,6 +442,8 @@ func (s *ClusteringCompactionTaskSuite) TestProcessExecuting() {
 				},
 			},
 		}, nil).Once()
+		s.session.EXPECT().DropCompactionPlan(mock.Anything, mock.Anything).Return(nil).Once()
+
 		time.Sleep(time.Second * 1)
 		s.Equal(true, task.Process())
 		s.Equal(datapb.CompactionTaskState_cleaned, task.GetState())

diff --git a/internal/datanode/compaction/clustering_compactor.go b/internal/datanode/compaction/clustering_compactor.go
@@ -277,6 +277,8 @@ func (t *clusteringCompactionTask) Compact() (*datapb.CompactionPlanResult, erro
 		WithLabelValues(fmt.Sprint(paramtable.GetNodeID()), t.plan.GetType().String()).
 		Observe(float64(t.tr.ElapseSpan().Milliseconds()))
 	log.Info("Clustering compaction finished", zap.Duration("elapse", t.tr.ElapseSpan()), zap.Int64("flushTimes", t.flushCount.Load()))
+	// clear the buffer cache
+	t.keyToBufferFunc = nil
 
 	return planResult, nil
 }