Merge pull request #14 from reddit/bug-fix-47072

Bug fix Segfault on ChannelLevelScoreBalancer use

Author: kumaramit01

internal/querycoordv2/observers/replica_observer.go

@@ -47,10 +47,11 @@ type ReplicaObserver struct {
 	stopOnce  sync.Once
 }
 
-func NewReplicaObserver(meta *meta.Meta, distMgr *meta.DistributionManager) *ReplicaObserver {
+func NewReplicaObserver(meta *meta.Meta, distMgr *meta.DistributionManager, targetMgr meta.TargetManagerInterface) *ReplicaObserver {
 	return &ReplicaObserver{
-		meta:    meta,
-		distMgr: distMgr,
+		meta:      meta,
+		distMgr:   distMgr,
+		targetMgr: targetMgr,
 	}
 }
 
@@ -179,6 +180,10 @@ func (ob *ReplicaObserver) checkNodesInReplica() {
 		replicas := ob.meta.ReplicaManager.GetByCollection(ctx, collectionID)
 		for _, replica := range replicas {
 			if enableChannelExclusiveMode && !replica.IsChannelExclusiveModeEnabled() {
+				if ob.targetMgr == nil {
+					log.Warn("targetMgr is nil, cannot enable channel exclusive mode")
+					continue
+				}
 				// register channel for enable exclusive mode
 				mutableReplica := replica.CopyForWrite()
 				channels := ob.targetMgr.GetDmChannelsByCollection(ctx, collectionID, meta.CurrentTargetFirst)

internal/querycoordv2/observers/replica_observer_test.go

@@ -30,7 +30,6 @@ import (
 	"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
 	"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_balancer"
 	"github.com/milvus-io/milvus/internal/querycoordv2/meta"
-	. "github.com/milvus-io/milvus/internal/querycoordv2/params"
 	"github.com/milvus-io/milvus/internal/querycoordv2/session"
 	"github.com/milvus-io/milvus/internal/querycoordv2/utils"
 	"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer"
@@ -43,15 +42,19 @@ import (
 	"github.com/milvus-io/milvus/pkg/v2/util/etcd"
 	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
+
+	. "github.com/milvus-io/milvus/internal/querycoordv2/params"
 )
 
 type ReplicaObserverSuite struct {
 	suite.Suite
 
 	kv kv.MetaKv
 	// dependency
-	meta    *meta.Meta
-	distMgr *meta.DistributionManager
+	meta      *meta.Meta
+	distMgr   *meta.DistributionManager
+	targetMgr meta.TargetManagerInterface
+	broker    *meta.MockBroker
 
 	nodeMgr  *session.NodeManager
 	observer *ReplicaObserver
@@ -89,9 +92,11 @@ func (suite *ReplicaObserverSuite) SetupTest() {
 	idAllocator := RandomIncrementIDAllocator()
 	suite.nodeMgr = session.NewNodeManager()
 	suite.meta = meta.NewMeta(idAllocator, store, suite.nodeMgr)
+	suite.broker = meta.NewMockBroker(suite.T())
+	suite.targetMgr = meta.NewTargetManager(suite.broker, suite.meta)
 
 	suite.distMgr = meta.NewDistributionManager(suite.nodeMgr)
-	suite.observer = NewReplicaObserver(suite.meta, suite.distMgr)
+	suite.observer = NewReplicaObserver(suite.meta, suite.distMgr, suite.targetMgr)
 	suite.observer.Start()
 	suite.collectionID = int64(1000)
 	suite.partitionID = int64(100)
@@ -244,7 +249,7 @@ func (suite *ReplicaObserverSuite) TestCheckSQnodesInReplica() {
 		}
 	})
 	balance.Register(b)
-	suite.observer = NewReplicaObserver(suite.meta, suite.distMgr)
+	suite.observer = NewReplicaObserver(suite.meta, suite.distMgr, suite.targetMgr)
 	suite.observer.Start()
 
 	ctx := context.Background()

internal/querycoordv2/server.go

@@ -453,6 +453,7 @@ func (s *Server) initObserver() {
 	s.replicaObserver = observers.NewReplicaObserver(
 		s.meta,
 		s.dist,
+		s.targetMgr,
 	)
 
 	s.resourceObserver = observers.NewResourceObserver(s.meta)

internal/querycoordv2/services_test.go

@@ -707,7 +707,7 @@ func (suite *ServiceSuite) TestTransferNode() {
 
 	server.resourceObserver = observers.NewResourceObserver(server.meta)
 	server.resourceObserver.Start()
-	server.replicaObserver = observers.NewReplicaObserver(server.meta, server.dist)
+	server.replicaObserver = observers.NewReplicaObserver(server.meta, server.dist, server.targetMgr)
 	server.replicaObserver.Start()
 	defer server.resourceObserver.Stop()
 	defer server.replicaObserver.Stop()
@@ -1917,6 +1917,7 @@ func (suite *ServiceSuite) TestHandleNodeUp() {
 	suite.server.replicaObserver = observers.NewReplicaObserver(
 		suite.server.meta,
 		suite.server.dist,
+		suite.server.targetMgr,
 	)
 	suite.server.resourceObserver = observers.NewResourceObserver(
 		suite.server.meta,