Merge remote-tracking branch 'upstream/master' into upgrade-golang-version

Author: XuanYang-cn

internal/core/src/segcore/SegmentSealedImpl.cpp

@@ -1591,7 +1591,11 @@ SegmentSealedImpl::bulk_subscript(
         return fill_with_empty(field_id, 0);
     }
 
-    auto column = fields_.at(field_id);
+    std::shared_ptr<SingleChunkColumnBase> column;
+    {
+        std::shared_lock lck(mutex_);
+        column = fields_.at(field_id);
+    }
     auto ret = fill_with_empty(field_id, count);
     if (column->IsNullable()) {
         auto dst = ret->mutable_valid_data()->mutable_data();

internal/proxy/impl.go

@@ -5439,10 +5439,6 @@ func (node *Proxy) validateOperatePrivilegeV2Params(req *milvuspb.OperatePrivile
 	if err := ValidatePrivilege(req.Grantor.Privilege.Name); err != nil {
 		return err
 	}
-	// validate built-in privilege group params
-	if err := ValidateBuiltInPrivilegeGroup(req.Grantor.Privilege.Name, req.DbName, req.CollectionName); err != nil {
-		return err
-	}
 	if req.Type != milvuspb.OperatePrivilegeType_Grant && req.Type != milvuspb.OperatePrivilegeType_Revoke {
 		return merr.WrapErrParameterInvalidMsg("the type in the request not grant or revoke")
 	}

internal/proxy/util.go

@@ -1120,31 +1120,6 @@ func ValidatePrivilege(entity string) error {
 	return validateName(entity, "Privilege")
 }
 
-func ValidateBuiltInPrivilegeGroup(entity string, dbName string, collectionName string) error {
-	if !util.IsBuiltinPrivilegeGroup(entity) {
-		return nil
-	}
-	switch {
-	case strings.HasPrefix(entity, milvuspb.PrivilegeLevel_Cluster.String()):
-		if !util.IsAnyWord(dbName) || !util.IsAnyWord(collectionName) {
-			return merr.WrapErrParameterInvalidMsg("dbName and collectionName should be * for the cluster level privilege: %s", entity)
-		}
-		return nil
-	case strings.HasPrefix(entity, milvuspb.PrivilegeLevel_Database.String()):
-		if collectionName != "" && collectionName != util.AnyWord {
-			return merr.WrapErrParameterInvalidMsg("collectionName should be * for the database level privilege: %s", entity)
-		}
-		return nil
-	case strings.HasPrefix(entity, milvuspb.PrivilegeLevel_Collection.String()):
-		if util.IsAnyWord(dbName) && !util.IsAnyWord(collectionName) && collectionName != "" {
-			return merr.WrapErrParameterInvalidMsg("please specify database name for the collection level privilege: %s", entity)
-		}
-		return nil
-	default:
-		return nil
-	}
-}
-
 func GetCurUserFromContext(ctx context.Context) (string, error) {
 	return contextutil.GetCurUserFromContext(ctx)
 }

internal/querycoordv2/balance/balance.go

@@ -46,10 +46,13 @@ func (segPlan *SegmentAssignPlan) String() string {
 }
 
 type ChannelAssignPlan struct {
-	Channel *meta.DmChannel
-	Replica *meta.Replica
-	From    int64
-	To      int64
+	Channel      *meta.DmChannel
+	Replica      *meta.Replica
+	From         int64
+	To           int64
+	FromScore    int64
+	ToScore      int64
+	ChannelScore int64
 }
 
 func (chanPlan *ChannelAssignPlan) String() string {

internal/querycoordv2/balance/rowcount_based_balancer.go

@@ -393,7 +393,7 @@ func newNodeItem(currentScore int, nodeID int64) nodeItem {
 
 func (b *nodeItem) getPriority() int {
 	// if node lacks more score between assignedScore and currentScore, then higher priority
-	return int(b.currentScore - b.assignedScore)
+	return int(math.Ceil(b.currentScore - b.assignedScore))
 }
 
 func (b *nodeItem) setPriority(priority int) {

internal/querycoordv2/balance/score_based_balancer.go

@@ -191,19 +191,19 @@ func (b *ScoreBasedBalancer) assignChannel(br *balanceReport, collectionID int64
 			}
 
 			from := int64(-1)
-			// fromScore := int64(0)
+			fromScore := int64(0)
 			if sourceNode != nil {
 				from = sourceNode.nodeID
-				// fromScore = int64(sourceNode.getPriority())
+				fromScore = int64(sourceNode.getPriority())
 			}
 
 			plan := ChannelAssignPlan{
-				From:    from,
-				To:      targetNode.nodeID,
-				Channel: ch,
-				// FromScore:    fromScore,
-				// ToScore:      int64(targetNode.getPriority()),
-				// SegmentScore: int64(scoreChanges),
+				From:         from,
+				To:           targetNode.nodeID,
+				Channel:      ch,
+				FromScore:    fromScore,
+				ToScore:      int64(targetNode.getPriority()),
+				ChannelScore: int64(scoreChanges),
 			}
 			br.AddRecord(StrRecordf("add segment plan %s", plan))
 			plans = append(plans, plan)
@@ -487,6 +487,20 @@ func (b *ScoreBasedBalancer) BalanceReplica(ctx context.Context, replica *meta.R
 	return segmentPlans, channelPlans
 }
 
+func (b *ScoreBasedBalancer) genStoppingChannelPlan(ctx context.Context, replica *meta.Replica, rwNodes []int64, roNodes []int64) []ChannelAssignPlan {
+	channelPlans := make([]ChannelAssignPlan, 0)
+	for _, nodeID := range roNodes {
+		dmChannels := b.dist.ChannelDistManager.GetByCollectionAndFilter(replica.GetCollectionID(), meta.WithNodeID2Channel(nodeID))
+		plans := b.AssignChannel(ctx, replica.GetCollectionID(), dmChannels, rwNodes, false)
+		for i := range plans {
+			plans[i].From = nodeID
+			plans[i].Replica = replica
+		}
+		channelPlans = append(channelPlans, plans...)
+	}
+	return channelPlans
+}
+
 func (b *ScoreBasedBalancer) genStoppingSegmentPlan(ctx context.Context, replica *meta.Replica, onlineNodes []int64, offlineNodes []int64) []SegmentAssignPlan {
 	segmentPlans := make([]SegmentAssignPlan, 0)
 	for _, nodeID := range offlineNodes {

internal/querycoordv2/balance/score_based_balancer_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
+	"go.uber.org/atomic"
 
 	etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
 	"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
@@ -1470,3 +1471,127 @@ func (suite *ScoreBasedBalancerTestSuite) TestBalanceChannelOnChannelExclusive()
 	_, channelPlans = suite.getCollectionBalancePlans(balancer, 3)
 	suite.Len(channelPlans, 2)
 }
+
+func (suite *ScoreBasedBalancerTestSuite) TestBalanceChannelOnStoppingNode() {
+	ctx := context.Background()
+	balancer := suite.balancer
+
+	// mock 10 collections with each collection has 1 channel
+	collectionNum := 10
+	channelNum := 1
+	for i := 1; i <= collectionNum; i++ {
+		collectionID := int64(i)
+		collection := utils.CreateTestCollection(collectionID, int32(1))
+		collection.LoadPercentage = 100
+		collection.Status = querypb.LoadStatus_Loaded
+		balancer.meta.CollectionManager.PutCollection(ctx, collection)
+		balancer.meta.CollectionManager.PutPartition(ctx, utils.CreateTestPartition(collectionID, collectionID))
+		balancer.meta.ReplicaManager.Spawn(ctx, collectionID, map[string]int{meta.DefaultResourceGroupName: 1}, nil)
+
+		channels := make([]*datapb.VchannelInfo, channelNum)
+		for i := 0; i < channelNum; i++ {
+			channels[i] = &datapb.VchannelInfo{CollectionID: collectionID, ChannelName: fmt.Sprintf("channel-%d-%d", collectionID, i)}
+		}
+		suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, collectionID).Return(
+			channels, nil, nil)
+		suite.broker.EXPECT().GetPartitions(mock.Anything, collectionID).Return([]int64{collectionID}, nil).Maybe()
+		balancer.targetMgr.UpdateCollectionNextTarget(ctx, collectionID)
+		balancer.targetMgr.UpdateCollectionCurrentTarget(ctx, collectionID)
+	}
+
+	// mock querynode-1 to node manager
+	nodeInfo := session.NewNodeInfo(session.ImmutableNodeInfo{
+		NodeID:   1,
+		Address:  "127.0.0.1:0",
+		Hostname: "localhost",
+		Version:  common.Version,
+	})
+	nodeInfo.SetState(session.NodeStateNormal)
+	suite.balancer.nodeManager.Add(nodeInfo)
+	suite.balancer.meta.ResourceManager.HandleNodeUp(ctx, 1)
+	utils.RecoverAllCollection(balancer.meta)
+
+	// mock channel distribution
+	channelDist := make([]*meta.DmChannel, 0)
+	for i := 1; i <= collectionNum; i++ {
+		collectionID := int64(i)
+		for i := 0; i < channelNum; i++ {
+			channelDist = append(channelDist, &meta.DmChannel{
+				VchannelInfo: &datapb.VchannelInfo{CollectionID: collectionID, ChannelName: fmt.Sprintf("channel-%d-%d", collectionID, i)}, Node: 1,
+			})
+		}
+	}
+	balancer.dist.ChannelDistManager.Update(1, channelDist...)
+
+	// assert balance channel won't happens on 1 querynode
+	ret := make([]ChannelAssignPlan, 0)
+	for i := 1; i <= collectionNum; i++ {
+		collectionID := int64(i)
+		_, channelPlans := suite.getCollectionBalancePlans(balancer, collectionID)
+		ret = append(ret, channelPlans...)
+	}
+	suite.Len(ret, 0)
+
+	// mock querynode-2 and querynode-3 to node manager
+	nodeInfo2 := session.NewNodeInfo(session.ImmutableNodeInfo{
+		NodeID:   2,
+		Address:  "127.0.0.1:0",
+		Hostname: "localhost",
+		Version:  common.Version,
+	})
+	suite.balancer.nodeManager.Add(nodeInfo2)
+	suite.balancer.meta.ResourceManager.HandleNodeUp(ctx, 2)
+	// mock querynode-2 and querynode-3 to node manager
+	nodeInfo3 := session.NewNodeInfo(session.ImmutableNodeInfo{
+		NodeID:   3,
+		Address:  "127.0.0.1:0",
+		Hostname: "localhost",
+		Version:  common.Version,
+	})
+	suite.balancer.nodeManager.Add(nodeInfo3)
+	suite.balancer.meta.ResourceManager.HandleNodeUp(ctx, 3)
+	utils.RecoverAllCollection(balancer.meta)
+	// mock querynode-1 to stopping, trigger stopping balance, expect to generate 10 balance channel task, and 5 for node-2, 5 for node-3
+	nodeInfo.SetState(session.NodeStateStopping)
+	suite.balancer.meta.ResourceManager.HandleNodeDown(ctx, 1)
+	utils.RecoverAllCollection(balancer.meta)
+
+	node2Counter := atomic.NewInt32(0)
+	node3Counter := atomic.NewInt32(0)
+
+	suite.mockScheduler.ExpectedCalls = nil
+	suite.mockScheduler.EXPECT().GetSegmentTaskDelta(mock.Anything, mock.Anything).Return(0).Maybe()
+	suite.mockScheduler.EXPECT().GetChannelTaskDelta(mock.Anything, mock.Anything).RunAndReturn(func(nodeID, collection int64) int {
+		if collection == -1 {
+			if nodeID == 2 {
+				return int(node2Counter.Load())
+			}
+
+			if nodeID == 3 {
+				return int(node3Counter.Load())
+			}
+		}
+		return 0
+	})
+	suite.mockScheduler.EXPECT().GetSegmentTaskNum(mock.Anything, mock.Anything).Return(0).Maybe()
+	suite.mockScheduler.EXPECT().GetChannelTaskNum(mock.Anything, mock.Anything).Return(0).Maybe()
+
+	for i := 1; i <= collectionNum; i++ {
+		collectionID := int64(i)
+		_, channelPlans := suite.getCollectionBalancePlans(balancer, collectionID)
+		suite.Len(channelPlans, 1)
+		if channelPlans[0].To == 2 {
+			node2Counter.Inc()
+		}
+
+		if channelPlans[0].To == 3 {
+			node3Counter.Inc()
+		}
+
+		if i%2 == 0 {
+			suite.Equal(node2Counter.Load(), node3Counter.Load())
+		}
+	}
+	suite.Equal(node2Counter.Load(), int32(5))
+	suite.Equal(node3Counter.Load(), int32(5))
+}

internal/rootcoord/meta_table.go

@@ -1525,7 +1525,7 @@ func (mt *MetaTable) RestoreRBAC(ctx context.Context, tenant string, meta *milvu
 	return mt.catalog.RestoreRBAC(ctx, tenant, meta)
 }
 
-// check if the privielge group name is defined by users
+// check if the privilege group name is defined by users
 func (mt *MetaTable) IsCustomPrivilegeGroup(ctx context.Context, groupName string) (bool, error) {
 	privGroups, err := mt.catalog.ListPrivilegeGroups(ctx)
 	if err != nil {
@@ -1641,7 +1641,7 @@ func (mt *MetaTable) OperatePrivilegeGroup(ctx context.Context, groupName string
 			if group.GroupName == p.Name {
 				privileges = append(privileges, group.Privileges...)
 			} else {
-				return merr.WrapErrParameterInvalidMsg("there is no privilege name or privielge group name [%s] defined in system to operate", p.Name)
+				return merr.WrapErrParameterInvalidMsg("there is no privilege name or privilege group name [%s] defined in system to operate", p.Name)
 			}
 		}
 	}

internal/rootcoord/rbac_task.go

@@ -179,7 +179,7 @@ func executeOperatePrivilegeTaskSteps(ctx context.Context, core *Core, in *milvu
 		}
 		grants := []*milvuspb.GrantEntity{in.Entity}
 
-		allGroups, err := core.getPrivilegeGroups(ctx)
+		allGroups, err := core.getDefaultAndCustomPrivilegeGroups(ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -275,12 +275,12 @@ func executeOperatePrivilegeGroupTaskSteps(ctx context.Context, core *Core, in *
 					return p.Name
 				})
 
-				// check if privileges are the same object type
-				objectTypes := lo.SliceToMap(newPrivs, func(p *milvuspb.PrivilegeEntity) (string, struct{}) {
-					return util.GetObjectType(p.Name), struct{}{}
+				// check if privileges are the same privilege level
+				privilegeLevels := lo.SliceToMap(newPrivs, func(p *milvuspb.PrivilegeEntity) (string, struct{}) {
+					return util.GetPrivilegeLevel(p.Name), struct{}{}
 				})
-				if len(objectTypes) > 1 {
-					return nil, errors.New("privileges are not the same object type")
+				if len(privilegeLevels) > 1 {
+					return nil, errors.New("privileges are not the same privilege level")
 				}
 			case milvuspb.OperatePrivilegeGroupType_RemovePrivilegesFromGroup:
 				newPrivs, _ := lo.Difference(v, in.Privileges)

internal/rootcoord/root_coord.go

@@ -2595,6 +2595,10 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
 			ctxLog.Error("", zap.Error(err))
 			return merr.StatusWithErrorCode(err, commonpb.ErrorCode_OperatePrivilegeFailure), nil
 		}
+		if err := c.validatePrivilegeGroupParams(ctx, privName, in.Entity.DbName, in.Entity.ObjectName); err != nil {
+			ctxLog.Error("", zap.Error(err))
+			return merr.StatusWithErrorCode(err, commonpb.ErrorCode_OperatePrivilegeFailure), nil
+		}
 		// set up object type for metastore, to be compatible with v1 version
 		in.Entity.Object.Name = util.GetObjectType(privName)
 	default:
@@ -2656,6 +2660,42 @@ func (c *Core) operatePrivilegeCommonCheck(ctx context.Context, in *milvuspb.Ope
 	return nil
 }
 
+func (c *Core) validatePrivilegeGroupParams(ctx context.Context, entity string, dbName string, collectionName string) error {
+	allGroups, err := c.getDefaultAndCustomPrivilegeGroups(ctx)
+	if err != nil {
+		return err
+	}
+	groups := lo.SliceToMap(allGroups, func(group *milvuspb.PrivilegeGroupInfo) (string, []*milvuspb.PrivilegeEntity) {
+		return group.GroupName, group.Privileges
+	})
+	privs, exists := groups[entity]
+	if !exists || len(privs) == 0 {
+		// it is a privilege, no need to check with other params
+		return nil
+	}
+	// since all privileges are same level in a group, just check the first privilege
+	level := util.GetPrivilegeLevel(privs[0].GetName())
+	switch level {
+	case milvuspb.PrivilegeLevel_Cluster.String():
+		if !util.IsAnyWord(dbName) || !util.IsAnyWord(collectionName) {
+			return merr.WrapErrParameterInvalidMsg("dbName and collectionName should be * for the cluster level privilege: %s", entity)
+		}
+		return nil
+	case milvuspb.PrivilegeLevel_Database.String():
+		if collectionName != "" && collectionName != util.AnyWord {
+			return merr.WrapErrParameterInvalidMsg("collectionName should be * for the database level privilege: %s", entity)
+		}
+		return nil
+	case milvuspb.PrivilegeLevel_Collection.String():
+		if util.IsAnyWord(dbName) && !util.IsAnyWord(collectionName) && collectionName != "" {
+			return merr.WrapErrParameterInvalidMsg("please specify database name for the collection level privilege: %s", entity)
+		}
+		return nil
+	default:
+		return errors.New("not found the privilege level")
+	}
+}
+
 func (c *Core) getMetastorePrivilegeName(ctx context.Context, privName string) (string, error) {
 	// if it is built-in privilege, return the privilege name directly
 	if util.IsPrivilegeNameDefined(privName) {
@@ -2757,7 +2797,7 @@ func (c *Core) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest)
 		}, nil
 	}
 	// expand privilege groups and turn to policies
-	allGroups, err := c.getPrivilegeGroups(ctx)
+	allGroups, err := c.getDefaultAndCustomPrivilegeGroups(ctx)
 	if err != nil {
 		errMsg := "fail to get privilege groups"
 		ctxLog.Warn(errMsg, zap.Error(err))
@@ -3131,8 +3171,8 @@ func (c *Core) expandPrivilegeGroups(ctx context.Context, grants []*milvuspb.Gra
 	}), nil
 }
 
-// getPrivilegeGroups returns default privilege groups and user-defined privilege groups.
-func (c *Core) getPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
+// getDefaultAndCustomPrivilegeGroups returns default privilege groups and user-defined privilege groups.
+func (c *Core) getDefaultAndCustomPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
 	allGroups, err := c.meta.ListPrivilegeGroups(ctx)
 	allGroups = append(allGroups, Params.RbacConfig.GetDefaultPrivilegeGroups()...)
 	if err != nil {