Fix current target may be updated to an invalid target (#21742)

Signed-off-by: Nyah01 <yang.cen@zilliz.com>

internal/querycoordv2/observers/collection_observer.go

@@ -37,6 +37,7 @@ type CollectionObserver struct {
 	dist                  *meta.DistributionManager
 	meta                  *meta.Meta
 	targetMgr             *meta.TargetManager
+	targetObserver        *TargetObserver
 	collectionLoadedCount map[int64]int
 	partitionLoadedCount  map[int64]int
 
@@ -47,12 +48,14 @@ func NewCollectionObserver(
 	dist *meta.DistributionManager,
 	meta *meta.Meta,
 	targetMgr *meta.TargetManager,
+	targetObserver *TargetObserver,
 ) *CollectionObserver {
 	return &CollectionObserver{
 		stopCh:                make(chan struct{}),
 		dist:                  dist,
 		meta:                  meta,
 		targetMgr:             targetMgr,
+		targetObserver:        targetObserver,
 		collectionLoadedCount: make(map[int64]int),
 		partitionLoadedCount:  make(map[int64]int),
 	}
@@ -201,9 +204,8 @@ func (ob *CollectionObserver) observeCollectionLoadStatus(collection *meta.Colle
 		return
 	}
 	ob.collectionLoadedCount[collection.GetCollectionID()] = loadedCount
-	if updated.LoadPercentage == 100 {
+	if updated.LoadPercentage == 100 && ob.targetObserver.Check(updated.GetCollectionID()) {
 		delete(ob.collectionLoadedCount, collection.GetCollectionID())
-		ob.targetMgr.UpdateCollectionCurrentTarget(updated.CollectionID)
 		updated.Status = querypb.LoadStatus_Loaded
 		ob.meta.CollectionManager.UpdateCollection(updated)
 
@@ -265,9 +267,8 @@ func (ob *CollectionObserver) observePartitionLoadStatus(partition *meta.Partiti
 		return
 	}
 	ob.partitionLoadedCount[partition.GetPartitionID()] = loadedCount
-	if updated.LoadPercentage == 100 {
+	if updated.LoadPercentage == 100 && ob.targetObserver.Check(updated.GetCollectionID()) {
 		delete(ob.partitionLoadedCount, partition.GetPartitionID())
-		ob.targetMgr.UpdateCollectionCurrentTarget(partition.GetCollectionID(), partition.GetPartitionID())
 		updated.Status = querypb.LoadStatus_Loaded
 		ob.meta.CollectionManager.PutPartition(updated)
 

internal/querycoordv2/observers/collection_observer_test.go

@@ -56,9 +56,10 @@ type CollectionObserverSuite struct {
 	broker      *meta.MockBroker
 
 	// Dependencies
-	dist      *meta.DistributionManager
-	meta      *meta.Meta
-	targetMgr *meta.TargetManager
+	dist           *meta.DistributionManager
+	meta           *meta.Meta
+	targetMgr      *meta.TargetManager
+	targetObserver *TargetObserver
 
 	// Test object
 	ob *CollectionObserver
@@ -180,18 +181,30 @@ func (suite *CollectionObserverSuite) SetupTest() {
 	suite.meta = meta.NewMeta(suite.idAllocator, suite.store)
 	suite.broker = meta.NewMockBroker(suite.T())
 	suite.targetMgr = meta.NewTargetManager(suite.broker, suite.meta)
+	suite.targetObserver = NewTargetObserver(suite.meta,
+		suite.targetMgr,
+		suite.dist,
+		suite.broker,
+	)
 
 	// Test object
 	suite.ob = NewCollectionObserver(
 		suite.dist,
 		suite.meta,
 		suite.targetMgr,
+		suite.targetObserver,
 	)
 
+	for _, collection := range suite.collections {
+		suite.broker.EXPECT().GetPartitions(mock.Anything, collection).Return(suite.partitions[collection], nil).Maybe()
+	}
+	suite.targetObserver.Start(context.Background())
+
 	suite.loadAll()
 }
 
 func (suite *CollectionObserverSuite) TearDownTest() {
+	suite.targetObserver.Stop()
 	suite.ob.Stop()
 	suite.kv.Close()
 }

internal/querycoordv2/observers/target_observer.go

@@ -30,6 +30,11 @@ import (
 	"github.com/milvus-io/milvus/internal/util/typeutil"
 )
 
+type checkRequest struct {
+	CollectionID int64
+	Notifier     chan bool
+}
+
 type targetUpdateRequest struct {
 	CollectionID  int64
 	Notifier      chan error
@@ -44,6 +49,7 @@ type TargetObserver struct {
 	distMgr   *meta.DistributionManager
 	broker    meta.Broker
 
+	manualCheck          chan checkRequest
 	nextTargetLastUpdate map[int64]time.Time
 	updateChan           chan targetUpdateRequest
 	mut                  sync.Mutex                // Guard readyNotifiers
@@ -59,6 +65,7 @@ func NewTargetObserver(meta *meta.Meta, targetMgr *meta.TargetManager, distMgr *
 		targetMgr:            targetMgr,
 		distMgr:              distMgr,
 		broker:               broker,
+		manualCheck:          make(chan checkRequest, 10),
 		nextTargetLastUpdate: make(map[int64]time.Time),
 		updateChan:           make(chan targetUpdateRequest),
 		readyNotifiers:       make(map[int64][]chan struct{}),
@@ -95,21 +102,48 @@ func (ob *TargetObserver) schedule(ctx context.Context) {
 			ob.clean()
 			ob.tryUpdateTarget()
 
-		case request := <-ob.updateChan:
-			err := ob.updateNextTarget(request.CollectionID)
+		case req := <-ob.manualCheck:
+			ob.check(req.CollectionID)
+			req.Notifier <- ob.targetMgr.IsCurrentTargetExist(req.CollectionID)
+
+		case req := <-ob.updateChan:
+			err := ob.updateNextTarget(req.CollectionID)
 			if err != nil {
-				close(request.ReadyNotifier)
+				close(req.ReadyNotifier)
 			} else {
 				ob.mut.Lock()
-				ob.readyNotifiers[request.CollectionID] = append(ob.readyNotifiers[request.CollectionID], request.ReadyNotifier)
+				ob.readyNotifiers[req.CollectionID] = append(ob.readyNotifiers[req.CollectionID], req.ReadyNotifier)
 				ob.mut.Unlock()
 			}
 
-			request.Notifier <- err
+			req.Notifier <- err
 		}
 	}
 }
 
+// Check checks whether the next target is ready,
+// and updates the current target if it is,
+// returns true if current target is not nil
+func (ob *TargetObserver) Check(collectionID int64) bool {
+	notifier := make(chan bool)
+	ob.manualCheck <- checkRequest{
+		CollectionID: collectionID,
+		Notifier:     notifier,
+	}
+	return <-notifier
+}
+
+func (ob *TargetObserver) check(collectionID int64) {
+	if ob.shouldUpdateCurrentTarget(collectionID) {
+		ob.updateCurrentTarget(collectionID)
+	}
+
+	if ob.shouldUpdateNextTarget(collectionID) {
+		// update next target in collection level
+		ob.updateNextTarget(collectionID)
+	}
+}
+
 // UpdateNextTarget updates the next target,
 // returns a channel which will be closed when the next target is ready,
 // or returns error if failed to pull target
@@ -138,14 +172,7 @@ func (ob *TargetObserver) ReleaseCollection(collectionID int64) {
 func (ob *TargetObserver) tryUpdateTarget() {
 	collections := ob.meta.GetAll()
 	for _, collectionID := range collections {
-		if ob.shouldUpdateCurrentTarget(collectionID) {
-			ob.updateCurrentTarget(collectionID)
-		}
-
-		if ob.shouldUpdateNextTarget(collectionID) {
-			// update next target in collection level
-			ob.updateNextTarget(collectionID)
-		}
+		ob.check(collectionID)
 	}
 
 	collectionSet := typeutil.NewUniqueSet(collections...)
@@ -199,12 +226,6 @@ func (ob *TargetObserver) updateNextTargetTimestamp(collectionID int64) {
 }
 
 func (ob *TargetObserver) shouldUpdateCurrentTarget(collectionID int64) bool {
-	// Collection observer will update the current target as loading done,
-	// avoid double updating, which will cause update current target to a unfinished next target
-	if !ob.targetMgr.IsCurrentTargetExist(collectionID) {
-		return false
-	}
-
 	replicaNum := ob.meta.CollectionManager.GetReplicaNumber(collectionID)
 
 	// check channel first

internal/querycoordv2/server.go

@@ -279,11 +279,6 @@ func (s *Server) initMeta() error {
 
 func (s *Server) initObserver() {
 	log.Info("init observers")
-	s.collectionObserver = observers.NewCollectionObserver(
-		s.dist,
-		s.meta,
-		s.targetMgr,
-	)
 	s.leaderObserver = observers.NewLeaderObserver(
 		s.dist,
 		s.meta,
@@ -296,6 +291,12 @@ func (s *Server) initObserver() {
 		s.dist,
 		s.broker,
 	)
+	s.collectionObserver = observers.NewCollectionObserver(
+		s.dist,
+		s.meta,
+		s.targetMgr,
+		s.targetObserver,
+	)
 }
 
 func (s *Server) afterStart() {

internal/querycoordv2/server_test.go

@@ -35,6 +35,7 @@ import (
 	"github.com/milvus-io/milvus/internal/querycoordv2/dist"
 	"github.com/milvus-io/milvus/internal/querycoordv2/meta"
 	"github.com/milvus-io/milvus/internal/querycoordv2/mocks"
+	"github.com/milvus-io/milvus/internal/querycoordv2/observers"
 	"github.com/milvus-io/milvus/internal/querycoordv2/params"
 	"github.com/milvus-io/milvus/internal/querycoordv2/session"
 	"github.com/milvus-io/milvus/internal/querycoordv2/task"
@@ -401,6 +402,18 @@ func (suite *ServerSuite) hackServer() {
 		suite.server.balancer,
 		suite.server.taskScheduler,
 	)
+	suite.server.targetObserver = observers.NewTargetObserver(
+		suite.server.meta,
+		suite.server.targetMgr,
+		suite.server.dist,
+		suite.broker,
+	)
+	suite.server.collectionObserver = observers.NewCollectionObserver(
+		suite.server.dist,
+		suite.server.meta,
+		suite.server.targetMgr,
+		suite.server.targetObserver,
+	)
 
 	suite.broker.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything).Return(&schemapb.CollectionSchema{}, nil).Maybe()
 	for _, collection := range suite.collections {