mcs: add set handler for balancer and alloc node for default keyspace group

pkg/keyspace/tso_keyspace_group.go

@@ -36,9 +36,10 @@
 )
 
 const (
-	defaultBalancerPolicy = balancer.PolicyRoundRobin
-	allocNodeTimeout      = 1 * time.Second
-	allocNodeInterval     = 10 * time.Millisecond
+	defaultBalancerPolicy                     = balancer.PolicyRoundRobin
+	allocNodesForDefaultKeyspaceGroupInterval = 1 * time.Second
+	allocNodesTimeout                         = 1 * time.Second
+	allocNodesInterval                        = 10 * time.Millisecond
 	// TODO: move it to etcdutil
 	watchEtcdChangeRetryInterval = 1 * time.Second
 	maxRetryTimes                = 25
@@ -114,6 +115,15 @@
 
 	m.Lock()
 	defer m.Unlock()
+
+	// If the etcd client is not nil, start the watch loop.
+	if m.client != nil {
+		m.nodesBalancer = balancer.GenByPolicy[string](m.policy)
+		m.serviceRegistryMap = make(map[string]string)
+		m.wg.Add(1)
+		go m.startWatchLoop()
+	}
+
 	// Ignore the error if default keyspace group already exists in the storage (e.g. PD restart/recover).
 	err := m.saveKeyspaceGroups([]*endpoint.KeyspaceGroup{defaultKeyspaceGroup}, false)
 	if err != nil && err != ErrKeyspaceGroupExists {
@@ -126,17 +136,16 @@
 		return err
 	}
 	for _, group := range groups {
+		if group.ID == utils.DefaultKeyspaceGroupID {
+			if len(group.Members) == 0 && m.client != nil {
+				// The default keyspace group should have one replica at least.
+				m.wg.Add(1)
+				go m.allocNodesForDefaultKeyspaceGroup()
+			}
+		}
 		userKind := endpoint.StringUserKind(group.UserKind)
 		m.groups[userKind].Put(group)
 	}
-
-	// If the etcd client is not nil, start the watch loop.
-	if m.client != nil {
-		m.nodesBalancer = balancer.GenByPolicy[string](m.policy)
-		m.serviceRegistryMap = make(map[string]string)
-		m.wg.Add(1)
-		go m.startWatchLoop()
-	}
 	return nil
 }
 
@@ -146,6 +155,31 @@
 	m.wg.Wait()
 }
 
+func (m *GroupManager) allocNodesForDefaultKeyspaceGroup() {
+	defer logutil.LogPanic()
+	defer m.wg.Done()
+	ticker := time.NewTicker(allocNodesForDefaultKeyspaceGroupInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-m.ctx.Done():
+			return
+		case <-ticker.C:
+		}
+		kg, err := m.GetKeyspaceGroupByID(utils.DefaultKeyspaceGroupID)
+		nodes := m.nodesBalancer.GetAll()
+		if err == nil && kg != nil && equalMembers(kg.Members, nodes) {
+			continue
+		}
+		err = m.SetNodesForKeyspaceGroup(utils.DefaultKeyspaceGroupID, nodes)
+		if err == nil {
+			log.Info("alloc nodes for default keyspace group", zap.Reflect("nodes", nodes))
+		} else {
+			log.Warn("failed to alloc nodes for default keyspace group", zap.Error(err))
+		}
+	}
+}
+
 func (m *GroupManager) startWatchLoop() {
 	defer logutil.LogPanic()
 	defer m.wg.Done()
@@ -156,12 +190,9 @@
 		revision int64
 		err      error
 	)
+	ticker := time.NewTicker(retryInterval)
+	defer ticker.Stop()
 	for i := 0; i < maxRetryTimes; i++ {
-		select {
-		case <-ctx.Done():
-			return
-		case <-time.After(retryInterval):
-		}
 		resp, err = etcdutil.EtcdKVGet(m.client, m.tsoServiceKey, clientv3.WithRange(m.tsoServiceEndKey))
 		if err == nil {
 			revision = resp.Header.Revision
@@ -177,6 +208,11 @@
 			break
 		}
 		log.Warn("failed to get tso service addrs from etcd and will retry", zap.Error(err))
+		select {
+		case <-m.ctx.Done():
+			return
+		case <-ticker.C:
+		}
 	}
 	if err != nil || revision == 0 {
 		log.Warn("failed to get tso service addrs from etcd finally when loading", zap.Error(err))
@@ -603,18 +639,23 @@
 	return nil
 }
 
-// GetNodesNum returns the number of nodes.
-func (m *GroupManager) GetNodesNum() int {
+// GetNodesCount returns the count of nodes.
+func (m *GroupManager) GetNodesCount() int {
+	if m.nodesBalancer == nil {
+		return 0
+	}
 	return m.nodesBalancer.Len()
 }
 
 // AllocNodesForKeyspaceGroup allocates nodes for the keyspace group.
-func (m *GroupManager) AllocNodesForKeyspaceGroup(id uint32, replica int) ([]endpoint.KeyspaceGroupMember, error) {
-	ctx, cancel := context.WithTimeout(m.ctx, allocNodeTimeout)
+func (m *GroupManager) AllocNodesForKeyspaceGroup(id uint32, desiredReplicaCount int) ([]endpoint.KeyspaceGroupMember, error) {
+	m.Lock()
+	defer m.Unlock()
+	ctx, cancel := context.WithTimeout(m.ctx, allocNodesTimeout)
 	defer cancel()
-	ticker := time.NewTicker(allocNodeInterval)
+	ticker := time.NewTicker(allocNodesInterval)
 	defer ticker.Stop()
-	nodes := make([]endpoint.KeyspaceGroupMember, 0, replica)
+	nodes := make([]endpoint.KeyspaceGroupMember, 0, desiredReplicaCount)
 	err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) error {
 		kg, err := m.store.LoadKeyspaceGroup(txn, id)
 		if err != nil {
@@ -628,14 +669,17 @@
 			exists[member.Address] = struct{}{}
 			nodes = append(nodes, member)
 		}
-		for len(exists) < replica {
+		if len(exists) >= desiredReplicaCount {
+			return nil
+		}
+		for len(exists) < desiredReplicaCount {
 			select {
 			case <-ctx.Done():
 				return nil
 			case <-ticker.C:
 			}
-			num := m.GetNodesNum()
-			if num < replica || num == 0 { // double check
+			countOfNodes := m.GetNodesCount()
+			if countOfNodes < desiredReplicaCount || countOfNodes == 0 { // double check
 				return ErrNoAvailableNode
 			}
 			addr := m.nodesBalancer.Next()
@@ -656,3 +700,51 @@
 	}
 	return nodes, nil
 }
+
+// SetNodesForKeyspaceGroup sets the nodes for the keyspace group.
+func (m *GroupManager) SetNodesForKeyspaceGroup(id uint32, nodes []string) error {
+	m.Lock()
+	defer m.Unlock()
+	return m.store.RunInTxn(m.ctx, func(txn kv.Txn) error {
+		kg, err := m.store.LoadKeyspaceGroup(txn, id)
+		if err != nil {
+			return err
+		}
+		if kg == nil {
+			return ErrKeyspaceGroupNotExists
+		}
+		members := make([]endpoint.KeyspaceGroupMember, 0, len(nodes))
+		for _, node := range nodes {
+			members = append(members, endpoint.KeyspaceGroupMember{Address: node})
+		}
+		kg.Members = members
+		return m.store.SaveKeyspaceGroup(txn, kg)
+	})
+}
+
+// IsExistNode checks if the node exists.
+func (m *GroupManager) IsExistNode(addr string) bool {
+	nodes := m.nodesBalancer.GetAll()
+	for _, node := range nodes {
+		if node == addr {
+			return true
+		}
+	}
+	return false
+}
+
+func equalMembers(a []endpoint.KeyspaceGroupMember, b []string) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	isExist := make(map[string]struct{}, len(b))
+	for _, node := range b {
+		isExist[node] = struct{}{}
+	}
+	for _, member := range a {
+		if _, ok := isExist[member.Address]; !ok {
+			return false
+		}
+	}
+	return true
+}

pkg/mcs/tso/server/server.go

@@ -219,6 +219,9 @@
 		log.Error("failed to get election member", errs.ZapError(err))
 		return false
 	}
+	if member == nil {
+		return false
+	}
 	return member.IsLeader()
 }
 

pkg/tso/keyspace_group_manager.go

@@ -728,6 +728,9 @@
 	if err != nil {
 		return nil, err
 	}
+	if am == nil {
+		return nil, nil
+	}
 	return am.GetMember(), nil
 }
 

server/apiv2/handlers/tso_keyspace_group.go

@@ -35,7 +35,8 @@
 	router.GET("", GetKeyspaceGroups)
 	router.GET("/:id", GetKeyspaceGroupByID)
 	router.DELETE("/:id", DeleteKeyspaceGroupByID)
-	router.POST("/:id/alloc", AllocNodeForKeyspaceGroup)
+	router.POST("/:id/alloc", AllocNodesForKeyspaceGroup)
+	router.POST("/:id/nodes", SetNodesForKeyspaceGroup)
 	router.POST("/:id/split", SplitKeyspaceGroupByID)
 	router.DELETE("/:id/split", FinishSplitKeyspaceByID)
 }
@@ -190,27 +191,27 @@
 	c.JSON(http.StatusOK, nil)
 }
 
-// AllocNodeForKeyspaceGroupParams defines the params for allocating nodes for keyspace groups.
-type AllocNodeForKeyspaceGroupParams struct {
+// AllocNodesForKeyspaceGroupParams defines the params for allocating nodes for keyspace groups.
+type AllocNodesForKeyspaceGroupParams struct {
 	Replica int `json:"replica"`
 }
 
-// AllocNodeForKeyspaceGroup allocates nodes for keyspace group.
-func AllocNodeForKeyspaceGroup(c *gin.Context) {
+// AllocNodesForKeyspaceGroup allocates nodes for keyspace group.
+func AllocNodesForKeyspaceGroup(c *gin.Context) {
 	id, err := validateKeyspaceGroupID(c)
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusBadRequest, "invalid keyspace group id")
 		return
 	}
 	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
 	manager := svr.GetKeyspaceGroupManager()
-	allocParams := &AllocNodeForKeyspaceGroupParams{}
+	allocParams := &AllocNodesForKeyspaceGroupParams{}
 	err = c.BindJSON(allocParams)
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusBadRequest, errs.ErrBindJSON.Wrap(err).GenWithStackByCause())
 		return
 	}
-	if manager.GetNodesNum() < allocParams.Replica || allocParams.Replica < 1 {
+	if manager.GetNodesCount() < allocParams.Replica || allocParams.Replica < 1 {
 		c.AbortWithStatusJSON(http.StatusBadRequest, "invalid replica, should be in [1, nodes_num]")
 		return
 	}
@@ -232,6 +233,54 @@
 	c.JSON(http.StatusOK, nodes)
 }
 
+// SetNodesForKeyspaceGroupParams defines the params for setting nodes for keyspace groups.
+// Notes: it should be used carefully.
+type SetNodesForKeyspaceGroupParams struct {
+	Nodes []string `json:"nodes"`
+}
+
+// SetNodesForKeyspaceGroup sets nodes for keyspace group.
+func SetNodesForKeyspaceGroup(c *gin.Context) {
+	id, err := validateKeyspaceGroupID(c)
+	if err != nil {
+		c.AbortWithStatusJSON(http.StatusBadRequest, "invalid keyspace group id")
+		return
+	}
+	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
+	manager := svr.GetKeyspaceGroupManager()
+	setParams := &SetNodesForKeyspaceGroupParams{}
+	err = c.BindJSON(setParams)
+	if err != nil {
+		c.AbortWithStatusJSON(http.StatusBadRequest, errs.ErrBindJSON.Wrap(err).GenWithStackByCause())
+		return
+	}
+	// check keyspace group whether exist
+	keyspaceGroup, err := manager.GetKeyspaceGroupByID(id)
+	if err != nil || keyspaceGroup == nil {
+		c.AbortWithStatusJSON(http.StatusBadRequest, "keyspace group does not exist")
+		return
+	}
+	// check nodes whether empty
+	if len(setParams.Nodes) == 0 {
+		c.AbortWithStatusJSON(http.StatusBadRequest, "invalid empty nodes")
+		return
+	}
+	// check nodes whether exist
+	for _, node := range setParams.Nodes {
+		if !manager.IsExistNode(node) {
+			c.AbortWithStatusJSON(http.StatusBadRequest, "node does not exist")
+			return
+		}
+	}
+	// set nodes
+	err = manager.SetNodesForKeyspaceGroup(id, setParams.Nodes)
+	if err != nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
+		return
+	}
+	c.JSON(http.StatusOK, nil)
+}
+
 func validateKeyspaceGroupID(c *gin.Context) (uint32, error) {
 	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
 	if err != nil {