fix: sync task still running after DataNode has stopped (milvus-io#38377)

jaime0815 · web-flow · commit 29e620fa6d5b · 2024-12-17T18:06:44.000+08:00
issue: milvus-io#38319 Signed-off-by: jaime <yun.zhang@zilliz.com>
diff --git a/configs/milvus.yaml b/configs/milvus.yaml
@@ -878,6 +878,8 @@ common:
   useVectorAsClusteringKey: false # if true, do clustering compaction and segment prune on vector field
   enableVectorClusteringKey: false # if true, enable vector clustering key and vector clustering compaction
   localRPCEnabled: false # enable local rpc for internal communication when mix or standalone mode.
+  sync:
+    taskPoolReleaseTimeoutSeconds: 60 # The maximum time to wait for the task to finish and release resources in the pool
 
 # QuotaConfig, configurations of Milvus quota and limits.
 # By default, we enable:
diff --git a/internal/datanode/data_node.go b/internal/datanode/data_node.go
@@ -402,6 +402,13 @@ func (node *DataNode) Stop() error {
 			node.writeBufferManager.Stop()
 		}
 
+		if node.syncMgr != nil {
+			err := node.syncMgr.Close()
+			if err != nil {
+				log.Error("sync manager close failed", zap.Error(err))
+			}
+		}
+
 		if node.allocator != nil {
 			log.Ctx(node.ctx).Info("close id allocator", zap.String("role", typeutil.DataNodeRole))
 			node.allocator.Close()
diff --git a/internal/datanode/importv2/scheduler_test.go b/internal/datanode/importv2/scheduler_test.go
@@ -246,11 +246,11 @@ func (s *SchedulerSuite) TestScheduler_Start_Import() {
 	cm.EXPECT().Reader(mock.Anything, mock.Anything).Return(&mockReader{Reader: ioReader}, nil)
 	s.cm = cm
 
-	s.syncMgr.EXPECT().SyncData(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, task syncmgr.Task, callbacks ...func(error) error) *conc.Future[struct{}] {
+	s.syncMgr.EXPECT().SyncData(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, task syncmgr.Task, callbacks ...func(error) error) (*conc.Future[struct{}], error) {
 		future := conc.Go(func() (struct{}, error) {
 			return struct{}{}, nil
 		})
-		return future
+		return future, nil
 	})
 	importReq := &datapb.ImportRequest{
 		JobID:        10,
@@ -307,11 +307,11 @@ func (s *SchedulerSuite) TestScheduler_Start_Import_Failed() {
 	cm.EXPECT().Reader(mock.Anything, mock.Anything).Return(&mockReader{Reader: ioReader}, nil)
 	s.cm = cm
 
-	s.syncMgr.EXPECT().SyncData(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, task syncmgr.Task, callbacks ...func(error) error) *conc.Future[struct{}] {
+	s.syncMgr.EXPECT().SyncData(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, task syncmgr.Task, callbacks ...func(error) error) (*conc.Future[struct{}], error) {
 		future := conc.Go(func() (struct{}, error) {
 			return struct{}{}, errors.New("mock err")
 		})
-		return future
+		return future, nil
 	})
 	importReq := &datapb.ImportRequest{
 		JobID:        10,
@@ -384,11 +384,11 @@ func (s *SchedulerSuite) TestScheduler_ReadFileStat() {
 }
 
 func (s *SchedulerSuite) TestScheduler_ImportFile() {
-	s.syncMgr.EXPECT().SyncData(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, task syncmgr.Task, callbacks ...func(error) error) *conc.Future[struct{}] {
+	s.syncMgr.EXPECT().SyncData(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, task syncmgr.Task, callbacks ...func(error) error) (*conc.Future[struct{}], error) {
 		future := conc.Go(func() (struct{}, error) {
 			return struct{}{}, nil
 		})
-		return future
+		return future, nil
 	})
 	var once sync.Once
 	data, err := testutil.CreateInsertData(s.schema, s.numRows)
diff --git a/internal/datanode/importv2/task_import.go b/internal/datanode/importv2/task_import.go
@@ -228,7 +228,7 @@ func (t *ImportTask) importFile(reader importutilv2.Reader) error {
 }
 
 func (t *ImportTask) sync(hashedData HashedData) ([]*conc.Future[struct{}], []syncmgr.Task, error) {
-	log.Info("start to sync import data", WrapLogFields(t)...)
+	log.Ctx(context.TODO()).Info("start to sync import data", WrapLogFields(t)...)
 	futures := make([]*conc.Future[struct{}], 0)
 	syncTasks := make([]syncmgr.Task, 0)
 	for channelIdx, datas := range hashedData {
@@ -256,7 +256,11 @@ func (t *ImportTask) sync(hashedData HashedData) ([]*conc.Future[struct{}], []sy
 			if err != nil {
 				return nil, nil, err
 			}
-			future := t.syncMgr.SyncData(t.ctx, syncTask)
+			future, err := t.syncMgr.SyncData(t.ctx, syncTask)
+			if err != nil {
+				log.Ctx(context.TODO()).Error("sync data failed", WrapLogFields(t, zap.Error(err))...)
+				return nil, nil, err
+			}
 			futures = append(futures, future)
 			syncTasks = append(syncTasks, syncTask)
 		}
diff --git a/internal/datanode/importv2/task_l0_import.go b/internal/datanode/importv2/task_l0_import.go
@@ -213,7 +213,7 @@ func (t *L0ImportTask) importL0(reader binlog.L0Reader) error {
 }
 
 func (t *L0ImportTask) syncDelete(delData []*storage.DeleteData) ([]*conc.Future[struct{}], []syncmgr.Task, error) {
-	log.Info("start to sync l0 delete data", WrapLogFields(t)...)
+	log.Ctx(context.TODO()).Info("start to sync l0 delete data", WrapLogFields(t)...)
 	futures := make([]*conc.Future[struct{}], 0)
 	syncTasks := make([]syncmgr.Task, 0)
 	for channelIdx, data := range delData {
@@ -231,7 +231,11 @@ func (t *L0ImportTask) syncDelete(delData []*storage.DeleteData) ([]*conc.Future
 		if err != nil {
 			return nil, nil, err
 		}
-		future := t.syncMgr.SyncData(t.ctx, syncTask)
+		future, err := t.syncMgr.SyncData(t.ctx, syncTask)
+		if err != nil {
+			log.Ctx(context.TODO()).Error("failed to sync l0 delete data", WrapLogFields(t, zap.Error(err))...)
+			return nil, nil, err
+		}
 		futures = append(futures, future)
 		syncTasks = append(syncTasks, syncTask)
 	}
diff --git a/internal/datanode/importv2/task_l0_import_test.go b/internal/datanode/importv2/task_l0_import_test.go
@@ -132,7 +132,7 @@ func (s *L0ImportSuite) TestL0PreImport() {
 
 func (s *L0ImportSuite) TestL0Import() {
 	s.syncMgr.EXPECT().SyncData(mock.Anything, mock.Anything).
-		RunAndReturn(func(ctx context.Context, task syncmgr.Task, callbacks ...func(error) error) *conc.Future[struct{}] {
+		RunAndReturn(func(ctx context.Context, task syncmgr.Task, callbacks ...func(error) error) (*conc.Future[struct{}], error) {
 			alloc := allocator.NewMockAllocator(s.T())
 			alloc.EXPECT().Alloc(mock.Anything).Return(1, int64(s.delCnt)+1, nil)
 			task.(*syncmgr.SyncTask).WithAllocator(alloc)
@@ -147,7 +147,7 @@ func (s *L0ImportSuite) TestL0Import() {
 			future := conc.Go(func() (struct{}, error) {
 				return struct{}{}, nil
 			})
-			return future
+			return future, nil
 		})
 
 	req := &datapb.ImportRequest{
diff --git a/internal/flushcommon/syncmgr/mock_sync_manager.go b/internal/flushcommon/syncmgr/mock_sync_manager.go
diff --git a/internal/flushcommon/syncmgr/sync_manager.go b/internal/flushcommon/syncmgr/sync_manager.go
@@ -47,8 +47,10 @@ type SyncMeta struct {
 //go:generate mockery --name=SyncManager --structname=MockSyncManager --output=./  --filename=mock_sync_manager.go --with-expecter --inpackage
 type SyncManager interface {
 	// SyncData is the method to submit sync task.
-	SyncData(ctx context.Context, task Task, callbacks ...func(error) error) *conc.Future[struct{}]
+	SyncData(ctx context.Context, task Task, callbacks ...func(error) error) (*conc.Future[struct{}], error)
 
+	// Close waits for the task to finish and then shuts down the sync manager.
+	Close() error
 	TaskStatsJSON() string
 }
 
@@ -97,13 +99,17 @@ func (mgr *syncManager) resizeHandler(evt *config.Event) {
 	}
 }
 
-func (mgr *syncManager) SyncData(ctx context.Context, task Task, callbacks ...func(error) error) *conc.Future[struct{}] {
+func (mgr *syncManager) SyncData(ctx context.Context, task Task, callbacks ...func(error) error) (*conc.Future[struct{}], error) {
+	if mgr.workerPool.IsClosed() {
+		return nil, fmt.Errorf("sync manager is closed")
+	}
+
 	switch t := task.(type) {
 	case *SyncTask:
 		t.WithChunkManager(mgr.chunkManager)
 	}
 
-	return mgr.safeSubmitTask(ctx, task, callbacks...)
+	return mgr.safeSubmitTask(ctx, task, callbacks...), nil
 }
 
 // safeSubmitTask submits task to SyncManager
@@ -126,6 +132,7 @@ func (mgr *syncManager) submit(ctx context.Context, key int64, task Task, callba
 	}
 	callbacks = append([]func(error) error{handler}, callbacks...)
 	log.Info("sync mgr sumbit task with key", zap.Int64("key", key))
+
 	return mgr.Submit(ctx, key, task, callbacks...)
 }
 
@@ -142,3 +149,8 @@ func (mgr *syncManager) TaskStatsJSON() string {
 	}
 	return string(ret)
 }
+
+func (mgr *syncManager) Close() error {
+	timeout := paramtable.Get().CommonCfg.SyncTaskPoolReleaseTimeoutSeconds.GetAsDuration(time.Second)
+	return mgr.workerPool.ReleaseTimeout(timeout)
+}
diff --git a/internal/flushcommon/syncmgr/sync_manager_test.go b/internal/flushcommon/syncmgr/sync_manager_test.go
@@ -173,11 +173,22 @@ func (s *SyncManagerSuite) TestSubmit() {
 		Timestamp:   100,
 	})
 
-	f := manager.SyncData(context.Background(), task)
+	f, err := manager.SyncData(context.Background(), task)
+	s.NoError(err)
 	s.NotNil(f)
 
-	_, err := f.Await()
+	_, err = f.Await()
+	s.NoError(err)
+}
+
+func (s *SyncManagerSuite) TestClose() {
+	manager := NewSyncManager(s.chunkManager)
+	err := manager.Close()
 	s.NoError(err)
+
+	f, err := manager.SyncData(context.Background(), nil)
+	s.Error(err)
+	s.Nil(f)
 }
 
 func (s *SyncManagerSuite) TestCompacted() {
@@ -202,10 +213,11 @@ func (s *SyncManagerSuite) TestCompacted() {
 		Timestamp:   100,
 	})
 
-	f := manager.SyncData(context.Background(), task)
+	f, err := manager.SyncData(context.Background(), task)
+	s.NoError(err)
 	s.NotNil(f)
 
-	_, err := f.Await()
+	_, err = f.Await()
 	s.NoError(err)
 	s.EqualValues(1001, segmentID.Load())
 }
@@ -254,7 +266,7 @@ func (s *SyncManagerSuite) TestUnexpectedError() {
 	task.EXPECT().Run(mock.Anything).Return(merr.WrapErrServiceInternal("mocked")).Once()
 	task.EXPECT().HandleError(mock.Anything)
 
-	f := manager.SyncData(context.Background(), task)
+	f, _ := manager.SyncData(context.Background(), task)
 	_, err := f.Await()
 	s.Error(err)
 }
@@ -268,7 +280,7 @@ func (s *SyncManagerSuite) TestTargetUpdateSameID() {
 	task.EXPECT().Run(mock.Anything).Return(errors.New("mock err")).Once()
 	task.EXPECT().HandleError(mock.Anything)
 
-	f := manager.SyncData(context.Background(), task)
+	f, _ := manager.SyncData(context.Background(), task)
 	_, err := f.Await()
 	s.Error(err)
 }
diff --git a/internal/flushcommon/writebuffer/write_buffer.go b/internal/flushcommon/writebuffer/write_buffer.go
@@ -322,7 +322,7 @@ func (wb *writeBufferBase) syncSegments(ctx context.Context, segmentIDs []int64)
 			}
 		}
 
-		result = append(result, wb.syncMgr.SyncData(ctx, syncTask, func(err error) error {
+		future, err := wb.syncMgr.SyncData(ctx, syncTask, func(err error) error {
 			if wb.taskObserverCallback != nil {
 				wb.taskObserverCallback(syncTask, err)
 			}
@@ -342,7 +342,11 @@ func (wb *writeBufferBase) syncSegments(ctx context.Context, segmentIDs []int64)
 				}
 			}
 			return nil
-		}))
+		})
+		if err != nil {
+			log.Fatal("failed to sync data", zap.Int64("segmentID", segmentID), zap.Error(err))
+		}
+		result = append(result, future)
 	}
 	return result
 }
@@ -643,7 +647,7 @@ func (wb *writeBufferBase) Close(ctx context.Context, drop bool) {
 			t.WithDrop()
 		}
 
-		f := wb.syncMgr.SyncData(ctx, syncTask, func(err error) error {
+		f, err := wb.syncMgr.SyncData(ctx, syncTask, func(err error) error {
 			if wb.taskObserverCallback != nil {
 				wb.taskObserverCallback(syncTask, err)
 			}
@@ -656,6 +660,9 @@ func (wb *writeBufferBase) Close(ctx context.Context, drop bool) {
 			}
 			return nil
 		})
+		if err != nil {
+			log.Fatal("failed to sync segment", zap.Int64("segmentID", id), zap.Error(err))
+		}
 		futures = append(futures, f)
 	}
 
diff --git a/internal/flushcommon/writebuffer/write_buffer_test.go b/internal/flushcommon/writebuffer/write_buffer_test.go
@@ -317,7 +317,7 @@ func (s *WriteBufferSuite) TestEvictBuffer() {
 		serializer.EXPECT().EncodeBuffer(mock.Anything, mock.Anything).Return(syncmgr.NewSyncTask(), nil)
 		s.syncMgr.EXPECT().SyncData(mock.Anything, mock.Anything, mock.Anything).Return(conc.Go[struct{}](func() (struct{}, error) {
 			return struct{}{}, nil
-		}))
+		}), nil)
 		defer func() {
 			s.wb.mut.Lock()
 			defer s.wb.mut.Unlock()
diff --git a/pkg/util/conc/pool.go b/pkg/util/conc/pool.go
diff --git a/pkg/util/paramtable/component_param.go b/pkg/util/paramtable/component_param.go
diff --git a/pkg/util/paramtable/component_param_test.go b/pkg/util/paramtable/component_param_test.go

Original file line number	Diff line number	Diff line change
`@@ -47,8 +47,10 @@ type SyncMeta struct {`
`47`	`47`	`//go:generate mockery --name=SyncManager --structname=MockSyncManager --output=./ --filename=mock_sync_manager.go --with-expecter --inpackage`
`48`	`48`	`type SyncManager interface {`
`49`	`49`	`// SyncData is the method to submit sync task.`
`50`		`- SyncData(ctx context.Context, task Task, callbacks ...func(error) error) *conc.Future[struct{}]`
	`50`	`+ SyncData(ctx context.Context, task Task, callbacks ...func(error) error) (*conc.Future[struct{}], error)`
`51`	`51`
	`52`	`+ // Close waits for the task to finish and then shuts down the sync manager.`
	`53`	`+ Close() error`
`52`	`54`	`TaskStatsJSON() string`
`53`	`55`	`}`
`54`	`56`
`@@ -97,13 +99,17 @@ func (mgr syncManager) resizeHandler(evt config.Event) {`
`97`	`99`	`}`
`98`	`100`	`}`
`99`	`101`
`100`		`-func (mgr syncManager) SyncData(ctx context.Context, task Task, callbacks ...func(error) error) conc.Future[struct{}] {`
	`102`	`+func (mgr syncManager) SyncData(ctx context.Context, task Task, callbacks ...func(error) error) (conc.Future[struct{}], error) {`
	`103`	`+ if mgr.workerPool.IsClosed() {`
	`104`	`+ return nil, fmt.Errorf("sync manager is closed")`
	`105`	`+ }`
	`106`	`+`
`101`	`107`	`switch t := task.(type) {`
`102`	`108`	`case *SyncTask:`
`103`	`109`	`t.WithChunkManager(mgr.chunkManager)`
`104`	`110`	`}`
`105`	`111`
`106`		`- return mgr.safeSubmitTask(ctx, task, callbacks...)`
	`112`	`+ return mgr.safeSubmitTask(ctx, task, callbacks...), nil`
`107`	`113`	`}`
`108`	`114`
`109`	`115`	`// safeSubmitTask submits task to SyncManager`
`@@ -126,6 +132,7 @@ func (mgr *syncManager) submit(ctx context.Context, key int64, task Task, callba`
`126`	`132`	`}`
`127`	`133`	`callbacks = append([]func(error) error{handler}, callbacks...)`
`128`	`134`	`log.Info("sync mgr sumbit task with key", zap.Int64("key", key))`
	`135`	`+`
`129`	`136`	`return mgr.Submit(ctx, key, task, callbacks...)`
`130`	`137`	`}`
`131`	`138`
`@@ -142,3 +149,8 @@ func (mgr *syncManager) TaskStatsJSON() string {`
`142`	`149`	`}`
`143`	`150`	`return string(ret)`
`144`	`151`	`}`
	`152`	`+`
	`153`	`+func (mgr *syncManager) Close() error {`
	`154`	`+ timeout := paramtable.Get().CommonCfg.SyncTaskPoolReleaseTimeoutSeconds.GetAsDuration(time.Second)`
	`155`	`+ return mgr.workerPool.ReleaseTimeout(timeout)`
	`156`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -322,7 +322,7 @@ func (wb *writeBufferBase) syncSegments(ctx context.Context, segmentIDs []int64)`
`322`	`322`	`}`
`323`	`323`	`}`
`324`	`324`
`325`		`- result = append(result, wb.syncMgr.SyncData(ctx, syncTask, func(err error) error {`
	`325`	`+ future, err := wb.syncMgr.SyncData(ctx, syncTask, func(err error) error {`
`326`	`326`	`if wb.taskObserverCallback != nil {`
`327`	`327`	`wb.taskObserverCallback(syncTask, err)`
`328`	`328`	`}`
`@@ -342,7 +342,11 @@ func (wb *writeBufferBase) syncSegments(ctx context.Context, segmentIDs []int64)`
`342`	`342`	`}`
`343`	`343`	`}`
`344`	`344`	`return nil`
`345`		`- }))`
	`345`	`+ })`
	`346`	`+ if err != nil {`
	`347`	`+ log.Fatal("failed to sync data", zap.Int64("segmentID", segmentID), zap.Error(err))`
	`348`	`+ }`
	`349`	`+ result = append(result, future)`
`346`	`350`	`}`
`347`	`351`	`return result`
`348`	`352`	`}`
`@@ -643,7 +647,7 @@ func (wb *writeBufferBase) Close(ctx context.Context, drop bool) {`
`643`	`647`	`t.WithDrop()`
`644`	`648`	`}`
`645`	`649`
`646`		`- f := wb.syncMgr.SyncData(ctx, syncTask, func(err error) error {`
	`650`	`+ f, err := wb.syncMgr.SyncData(ctx, syncTask, func(err error) error {`
`647`	`651`	`if wb.taskObserverCallback != nil {`
`648`	`652`	`wb.taskObserverCallback(syncTask, err)`
`649`	`653`	`}`
`@@ -656,6 +660,9 @@ func (wb *writeBufferBase) Close(ctx context.Context, drop bool) {`
`656`	`660`	`}`
`657`	`661`	`return nil`
`658`	`662`	`})`
	`663`	`+ if err != nil {`
	`664`	`+ log.Fatal("failed to sync segment", zap.Int64("segmentID", id), zap.Error(err))`
	`665`	`+ }`
`659`	`666`	`futures = append(futures, f)`
`660`	`667`	`}`
`661`	`668`