Refine QueryCoord stopping (#27371)

Signed-off-by: yah01 <[email protected]>

internal/querycoordv2/checkers/controller.go

@@ -42,7 +42,7 @@ var (
 )
 
 type CheckerController struct {
-	stopCh         chan struct{}
+	cancel         context.CancelFunc
 	manualCheckChs map[string]chan struct{}
 	meta           *meta.Meta
 	dist           *meta.DistributionManager
@@ -87,7 +87,6 @@ func NewCheckerController(
 	}
 
 	return &CheckerController{
-		stopCh:         make(chan struct{}),
 		manualCheckChs: manualCheckChs,
 		meta:           meta,
 		dist:           dist,
@@ -98,9 +97,12 @@ func NewCheckerController(
 	}
 }
 
-func (controller *CheckerController) Start(ctx context.Context) {
+func (controller *CheckerController) Start() {
+	ctx, cancel := context.WithCancel(context.Background())
+	controller.cancel = cancel
+
 	for checkerType := range controller.checkers {
-		go controller.StartChecker(ctx, checkerType)
+		go controller.startChecker(ctx, checkerType)
 	}
 }
 
@@ -119,19 +121,14 @@ func getCheckerInterval(checkerType string) time.Duration {
 	}
 }
 
-func (controller *CheckerController) StartChecker(ctx context.Context, checkerType string) {
+func (controller *CheckerController) startChecker(ctx context.Context, checkerType string) {
 	interval := getCheckerInterval(checkerType)
 	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
 
 	for {
 		select {
 		case <-ctx.Done():
-			log.Info("Checker stopped due to context canceled",
-				zap.String("type", checkerType))
-			return
-
-		case <-controller.stopCh:
 			log.Info("Checker stopped",
 				zap.String("type", checkerType))
 			return
@@ -149,7 +146,9 @@ func (controller *CheckerController) StartChecker(ctx context.Context, checkerTy
 
 func (controller *CheckerController) Stop() {
 	controller.stopOnce.Do(func() {
-		close(controller.stopCh)
+		if controller.cancel != nil {
+			controller.cancel()
+		}
 	})
 }
 

internal/querycoordv2/checkers/controller_test.go

@@ -17,7 +17,6 @@
 package checkers
 
 import (
-	"context"
 	"testing"
 	"time"
 
@@ -127,8 +126,7 @@ func (suite *CheckerControllerSuite) TestBasic() {
 
 	suite.balancer.EXPECT().AssignSegment(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	suite.balancer.EXPECT().AssignChannel(mock.Anything, mock.Anything).Return(nil)
-	ctx := context.Background()
-	suite.controller.Start(ctx)
+	suite.controller.Start()
 	defer suite.controller.Stop()
 
 	suite.Eventually(func() bool {

internal/querycoordv2/dist/dist_controller.go

@@ -88,8 +88,9 @@ func (dc *ControllerImpl) SyncAll(ctx context.Context) {
 func (dc *ControllerImpl) Stop() {
 	dc.mu.Lock()
 	defer dc.mu.Unlock()
-	for _, h := range dc.handlers {
+	for nodeID, h := range dc.handlers {
 		h.stop()
+		delete(dc.handlers, nodeID)
 	}
 }
 

internal/querycoordv2/dist/dist_handler.go

@@ -246,6 +246,10 @@ func (dh *distHandler) stop() {
 	dh.stopOnce.Do(func() {
 		close(dh.c)
 		dh.wg.Wait()
+
+		// clear dist
+		dh.dist.ChannelDistManager.Update(dh.nodeID)
+		dh.dist.SegmentDistManager.Update(dh.nodeID)
 	})
 }
 

internal/querycoordv2/job/job_test.go

@@ -163,10 +163,10 @@ func (suite *JobSuite) SetupTest() {
 		suite.dist,
 		suite.broker,
 	)
-	suite.targetObserver.Start(context.Background())
+	suite.targetObserver.Start()
 	suite.scheduler = NewScheduler()
 
-	suite.scheduler.Start(context.Background())
+	suite.scheduler.Start()
 	meta.GlobalFailedLoadCache = meta.NewFailedLoadCache()
 
 	suite.nodeMgr.Add(session.NewNodeInfo(1000, "localhost"))

internal/querycoordv2/job/scheduler.go

@@ -37,7 +37,7 @@ const (
 type jobQueue chan Job
 
 type Scheduler struct {
-	stopCh chan struct{}
+	cancel context.CancelFunc
 	wg     sync.WaitGroup
 
 	processors *typeutil.ConcurrentSet[int64] // Collections of having processor
@@ -49,73 +49,64 @@ type Scheduler struct {
 
 func NewScheduler() *Scheduler {
 	return &Scheduler{
-		stopCh:     make(chan struct{}),
 		processors: typeutil.NewConcurrentSet[int64](),
 		queues:     make(map[int64]jobQueue),
 		waitQueue:  make(jobQueue, waitQueueCap),
 	}
 }
 
-func (scheduler *Scheduler) Start(ctx context.Context) {
-	scheduler.schedule(ctx)
+func (scheduler *Scheduler) Start() {
+	ctx, cancel := context.WithCancel(context.Background())
+	scheduler.cancel = cancel
+
+	scheduler.wg.Add(1)
+	go func() {
+		defer scheduler.wg.Done()
+		scheduler.schedule(ctx)
+	}()
 }
 
 func (scheduler *Scheduler) Stop() {
 	scheduler.stopOnce.Do(func() {
-		close(scheduler.stopCh)
+		if scheduler.cancel != nil {
+			scheduler.cancel()
+		}
 		scheduler.wg.Wait()
 	})
 }
 
 func (scheduler *Scheduler) schedule(ctx context.Context) {
-	scheduler.wg.Add(1)
-	go func() {
-		defer scheduler.wg.Done()
-		ticker := time.NewTicker(500 * time.Millisecond)
-		defer ticker.Stop()
-		for {
-			select {
-			case <-ctx.Done():
-				log.Info("JobManager stopped due to context canceled")
-				return
-
-			case <-scheduler.stopCh:
-				log.Info("JobManager stopped")
-				for _, queue := range scheduler.queues {
-					close(queue)
-				}
-				return
+	ticker := time.NewTicker(500 * time.Millisecond)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			log.Info("JobManager stopped")
+			for _, queue := range scheduler.queues {
+				close(queue)
+			}
+			return
 
-			case job := <-scheduler.waitQueue:
-				queue, ok := scheduler.queues[job.CollectionID()]
-				if !ok {
-					queue = make(jobQueue, collectionQueueCap)
-					scheduler.queues[job.CollectionID()] = queue
-				}
-				queue <- job
-				scheduler.startProcessor(job.CollectionID(), queue)
-
-			case <-ticker.C:
-				for collection, queue := range scheduler.queues {
-					if len(queue) > 0 {
-						scheduler.startProcessor(collection, queue)
-					} else {
-						// Release resource if no job for the collection
-						close(queue)
-						delete(scheduler.queues, collection)
-					}
+		case job := <-scheduler.waitQueue:
+			queue, ok := scheduler.queues[job.CollectionID()]
+			if !ok {
+				queue = make(jobQueue, collectionQueueCap)
+				scheduler.queues[job.CollectionID()] = queue
+			}
+			queue <- job
+			scheduler.startProcessor(job.CollectionID(), queue)
+
+		case <-ticker.C:
+			for collection, queue := range scheduler.queues {
+				if len(queue) > 0 {
+					scheduler.startProcessor(collection, queue)
+				} else {
+					// Release resource if no job for the collection
+					close(queue)
+					delete(scheduler.queues, collection)
 				}
 			}
 		}
-	}()
-}
-
-func (scheduler *Scheduler) isStopped() bool {
-	select {
-	case <-scheduler.stopCh:
-		return true
-	default:
-		return false
 	}
 }
 
@@ -124,9 +115,6 @@ func (scheduler *Scheduler) Add(job Job) {
 }
 
 func (scheduler *Scheduler) startProcessor(collection int64, queue jobQueue) {
-	if scheduler.isStopped() {
-		return
-	}
 	if !scheduler.processors.Insert(collection) {
 		return
 	}