pdms: support primary/transfer api for scheduling and tso

pkg/election/leadership.go

@@ -64,6 +64,11 @@ type Leadership struct {
 	leaderKey   string
 	leaderValue string
 
+	leaderWatch struct {
+		syncutil.RWMutex
+		val bool
+	}
+
 	keepAliveCtx            context.Context
 	keepAliveCancelFunc     context.CancelFunc
 	keepAliveCancelFuncLock syncutil.Mutex
@@ -72,6 +77,10 @@ type Leadership struct {
 	campaignTimes []time.Time
 }
 
+func (ls *Leadership) GetLeaderValue() string {
+	return ls.leaderValue
+}
+
 // NewLeadership creates a new Leadership.
 func NewLeadership(client *clientv3.Client, leaderKey, purpose string) *Leadership {
 	leadership := &Leadership{
@@ -113,6 +122,20 @@ func (ls *Leadership) GetLeaderKey() string {
 	return ls.leaderKey
 }
 
+// SetLeaderWatch sets the leader watch flag.
+func (ls *Leadership) SetLeaderWatch(val bool) {
+	ls.leaderWatch.Lock()
+	ls.leaderWatch.val = val
+	ls.leaderWatch.Unlock()
+}
+
+// GetLeaderWatch gets the leader watch flag.
+func (ls *Leadership) GetLeaderWatch() bool {
+	ls.leaderWatch.RLock()
+	defer ls.leaderWatch.RUnlock()
+	return ls.leaderWatch.val
+}
+
 // GetCampaignTimesNum is used to get the campaign times of the leader within `campaignTimesRecordTimeout`.
 func (ls *Leadership) GetCampaignTimesNum() int {
 	if ls == nil {
@@ -375,6 +398,12 @@ func (ls *Leadership) Watch(serverCtx context.Context, revision int64) {
 						zap.Int64("revision", wresp.Header.Revision), zap.String("leader-key", ls.leaderKey), zap.String("purpose", ls.purpose))
 					return
 				}
+				// only API update the leader key to transfer the leader will meet
+				if ev.Type == mvccpb.PUT && ls.GetLeaderWatch() {
+					log.Info("[LeaderWatch] current leadership is updated",
+						zap.Int64("revision", wresp.Header.Revision), zap.String("leader-key", ls.leaderKey), zap.String("purpose", ls.purpose))
+					return
+				}
 			}
 			revision = wresp.Header.Revision + 1
 		}
@@ -393,4 +422,5 @@ func (ls *Leadership) Reset() {
 	}
 	ls.keepAliveCancelFuncLock.Unlock()
 	ls.getLease().Close()
+	ls.SetLeaderWatch(false)
 }

pkg/mcs/discovery/discover.go

@@ -15,12 +15,15 @@
 package discovery
 
 import (
+	"math/rand"
 	"strconv"
+	"time"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/log"
 	"github.com/tikv/pd/pkg/errs"
 	"github.com/tikv/pd/pkg/mcs/utils"
+	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/storage/kv"
 	"github.com/tikv/pd/pkg/utils/etcdutil"
 	"go.etcd.io/etcd/clientv3"
@@ -45,14 +48,14 @@ func Discover(cli *clientv3.Client, clusterID, serviceName string) ([]string, er
 }
 
 // GetMSMembers returns all the members of the specified service name.
-func GetMSMembers(name string, client *clientv3.Client) ([]ServiceRegistryEntry, error) {
-	switch name {
+func GetMSMembers(serviceName string, client *clientv3.Client) ([]ServiceRegistryEntry, error) {
+	switch serviceName {
 	case utils.TSOServiceName, utils.SchedulingServiceName, utils.ResourceManagerServiceName:
 		clusterID, err := etcdutil.GetClusterID(client, utils.ClusterIDPath)
 		if err != nil {
 			return nil, err
 		}
-		servicePath := ServicePath(strconv.FormatUint(clusterID, 10), name)
+		servicePath := ServicePath(strconv.FormatUint(clusterID, 10), serviceName)
 		resps, err := kv.NewSlowLogTxn(client).Then(clientv3.OpGet(servicePath, clientv3.WithPrefix())).Commit()
 		if err != nil {
 			return nil, errs.ErrEtcdKVGet.Wrap(err).GenWithStackByCause()
@@ -75,5 +78,56 @@ func GetMSMembers(name string, client *clientv3.Client) ([]ServiceRegistryEntry,
 		return entries, nil
 	}
 
-	return nil, errors.Errorf("unknown service name %s", name)
+	return nil, errors.Errorf("unknown service name %s", serviceName)
+}
+
+func TransferPrimary(client *clientv3.Client, serviceName, oldPrimary, newPrimary string) error {
+	log.Info("transfer primary", zap.String("service", serviceName), zap.String("from", oldPrimary), zap.String("to", newPrimary))
+	entries, err := GetMSMembers(serviceName, client)
+	if err != nil {
+		return err
+	}
+
+	// Do nothing when I am the only member of cluster.
+	if len(entries) == 1 && newPrimary == "" {
+		return errors.New("no valid follower to transfer primary")
+	}
+
+	var primaryIDs []string
+	var memberValues []string
+	for _, member := range entries {
+		if (newPrimary == "" && member.ServiceAddr != oldPrimary) || (newPrimary != "" && member.ServiceAddr == newPrimary) {
+			primaryIDs = append(primaryIDs, member.ServiceAddr)
+			memberValues = append(memberValues, string(member.MemberValue))
+		}
+	}
+	if len(primaryIDs) == 0 {
+		return errors.New("no valid follower to transfer primary")
+	}
+
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	nextPrimaryID := r.Intn(len(primaryIDs))
+
+	clusterID, err := etcdutil.GetClusterID(client, utils.ClusterIDPath)
+	if err != nil {
+		return errors.Errorf("failed to get cluster ID: %v", err)
+	}
+
+	var primaryKey string
+	switch serviceName {
+	case utils.SchedulingServiceName:
+		primaryKey = endpoint.SchedulingPrimaryPath(clusterID)
+	case utils.TSOServiceName:
+		tsoRootPath := endpoint.TSOSvcRootPath(clusterID)
+		primaryKey = endpoint.KeyspaceGroupPrimaryPath(tsoRootPath, utils.DefaultKeyspaceGroupID)
+	}
+
+	// update primary key to notify old primary server.
+	putResp, err := kv.NewSlowLogTxn(client).
+		Then(clientv3.OpPut(primaryKey, memberValues[nextPrimaryID])).
+		Commit()
+	if err != nil || !putResp.Succeeded {
+		return errors.Errorf("failed to write primary flag for %s", serviceName)
+	}
+	return nil
 }

pkg/mcs/discovery/registry_entry.go

@@ -23,11 +23,13 @@ import (
 
 // ServiceRegistryEntry is the registry entry of a service
 type ServiceRegistryEntry struct {
+	Name           string `json:"name"`
 	ServiceAddr    string `json:"service-addr"`
 	Version        string `json:"version"`
 	GitHash        string `json:"git-hash"`
 	DeployPath     string `json:"deploy-path"`
 	StartTimestamp int64  `json:"start-timestamp"`
+	MemberValue    []byte `json:"member-value"`
 }
 
 // Serialize this service registry entry

pkg/mcs/scheduling/server/server.go

@@ -17,6 +17,7 @@ package server
 import (
 	"context"
 	"fmt"
+	"github.com/tikv/pd/pkg/utils/etcdutil"
 	"net/http"
 	"os"
 	"os/signal"
@@ -128,6 +129,10 @@ func (s *Server) GetBackendEndpoints() string {
 	return s.cfg.BackendEndpoints
 }
 
+func (s *Server) GetParticipant() *member.Participant {
+	return s.participant
+}
+
 // SetLogLevel sets log level.
 func (s *Server) SetLogLevel(level string) error {
 	if !logutil.IsLevelLegal(level) {
@@ -243,6 +248,17 @@ func (s *Server) primaryElectionLoop() {
 			log.Info("the scheduling primary has changed, try to re-campaign a primary")
 		}
 
+		// To make sure the expected leader(if exist) and primary are on the same server.
+		expectedPrimary := utils.GetExpectedPrimary(s.participant.GetLeaderPath(), s.GetClient())
+		if expectedPrimary != "" && expectedPrimary != s.participant.GetLeadership().GetLeaderValue() {
+			log.Info("skip campaigning of scheduling primary and check later",
+				zap.String("server-name", s.Name()),
+				zap.String("target-primary-id", expectedPrimary),
+				zap.Uint64("member-id", s.participant.ID()))
+			time.Sleep(200 * time.Millisecond)
+			continue
+		}
+
 		s.campaignLeader()
 	}
 }
@@ -290,6 +306,9 @@ func (s *Server) campaignLeader() {
 	member.ServiceMemberGauge.WithLabelValues(serviceName).Set(1)
 	log.Info("scheduling primary is ready to serve", zap.String("scheduling-primary-name", s.participant.Name()))
 
+	exitPrimary := make(chan struct{})
+	go s.primaryWatch(ctx, exitPrimary)
+
 	leaderTicker := time.NewTicker(utils.LeaderTickInterval)
 	defer leaderTicker.Stop()
 
@@ -304,6 +323,47 @@ func (s *Server) campaignLeader() {
 			// Server is closed and it should return nil.
 			log.Info("server is closed")
 			return
+		case <-exitPrimary:
+			log.Info("no longer a primary/leader because primary have been updated, the scheduling primary/leader will step down")
+			return
+		}
+	}
+}
+
+func (s *Server) primaryWatch(ctx context.Context, exitPrimary chan struct{}) {
+	_, revision, err := s.participant.GetPersistentLeader()
+	if err != nil {
+		log.Error("[primary] getting the leader meets error", errs.ZapError(err))
+		return
+	}
+	log.Info("[primary] start to watch the primary", zap.Stringer("scheduling-primary", s.participant.GetLeader()))
+	// Watch will keep looping and never return unless the primary has changed.
+	s.participant.GetLeadership().SetLeaderWatch(true)
+	s.participant.GetLeadership().Watch(s.serverLoopCtx, revision+1)
+	s.participant.GetLeadership().SetLeaderWatch(false)
+
+	// only API update primary will set the expected leader
+	// check leader key whether deleted
+	leaderRaw, err := etcdutil.GetValue(s.participant.Client(), s.participant.GetLeaderPath())
+	if err != nil {
+		log.Error("[primary] get primary key error", zap.Error(err))
+		return
+	}
+	if leaderRaw == nil {
+		log.Info("[primary] leader key is deleted, the primary will step down")
+		return
+	}
+
+	utils.SetExpectedPrimary(s.participant.Client(), s.participant.GetLeaderPath())
+
+	s.participant.UnsetLeader()
+	for {
+		select {
+		case <-ctx.Done():
+			log.Info("[primary] exit the primary watch loop")
+			return
+		case exitPrimary <- struct{}{}:
+			return
 		}
 	}
 }
@@ -425,6 +485,7 @@ func (s *Server) startServer() (err error) {
 		GitHash:        versioninfo.PDGitHash,
 		DeployPath:     deployPath,
 		StartTimestamp: s.StartTimestamp(),
+		Name:           s.Name(),
 	}
 	uniqueName := s.cfg.GetAdvertiseListenAddr()
 	uniqueID := memberutil.GenerateUniqueID(uniqueName)
@@ -436,6 +497,7 @@ func (s *Server) startServer() (err error) {
 		ListenUrls: []string{s.cfg.GetAdvertiseListenAddr()},
 	}
 	s.participant.InitInfo(p, endpoint.SchedulingSvcRootPath(s.clusterID), utils.PrimaryKey, "primary election")
+	s.serviceID.MemberValue = []byte(s.participant.MemberValue())
 
 	s.service = &Service{Server: s}
 	s.AddServiceReadyCallback(s.startCluster)

pkg/mcs/tso/server/server.go

@@ -380,13 +380,17 @@ func (s *Server) startServer() (err error) {
 		GitHash:        versioninfo.PDGitHash,
 		DeployPath:     deployPath,
 		StartTimestamp: s.StartTimestamp(),
+		Name:           s.Name(),
 	}
 	s.keyspaceGroupManager = tso.NewKeyspaceGroupManager(
 		s.serverLoopCtx, s.serviceID, s.GetClient(), s.GetHTTPClient(), s.cfg.AdvertiseListenAddr,
 		s.clusterID, legacySvcRootPath, tsoSvcRootPath, s.cfg)
 	if err := s.keyspaceGroupManager.Initialize(); err != nil {
 		return err
 	}
+	// Initialize the service ID with the member value of the primary of the default keyspace group.
+	memberValue, err := s.GetMember(utils.DefaultKeyspaceID, utils.DefaultKeyspaceGroupID)
+	s.serviceID.MemberValue = []byte(memberValue.MemberValue())
 
 	s.tsoProtoFactory = &tsoutil.TSOProtoFactory{}
 	s.service = &Service{Server: s}

pkg/mcs/utils/util.go

@@ -32,6 +32,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/soheilhy/cmux"
 	"github.com/tikv/pd/pkg/errs"
+	"github.com/tikv/pd/pkg/storage/kv"
 	"github.com/tikv/pd/pkg/utils/apiutil"
 	"github.com/tikv/pd/pkg/utils/apiutil/multiservicesapi"
 	"github.com/tikv/pd/pkg/utils/etcdutil"
@@ -51,6 +52,9 @@ const (
 	ClusterIDPath = "/pd/cluster_id"
 	// retryInterval is the interval to retry.
 	retryInterval = time.Second
+	// ExpectedPrimary is the path to store the expected primary
+	// ONLY SET VALUE BY API
+	ExpectedPrimary = "expected_primary"
 )
 
 // InitClusterID initializes the cluster ID.
@@ -70,6 +74,49 @@ func InitClusterID(ctx context.Context, client *clientv3.Client) (id uint64, err
 	return 0, errors.Errorf("failed to init cluster ID after retrying %d times", maxRetryTimes)
 }
 
+// GetExpectedPrimary indicates API has changed the primary, ONLY SET VALUE BY API.
+func GetExpectedPrimary(keyPath string, client *clientv3.Client) string {
+	leader, err := etcdutil.GetValue(client, strings.Join([]string{keyPath, ExpectedPrimary}, "/"))
+	if err != nil {
+		log.Error("get expected primary key error", errs.ZapError(err))
+		return ""
+	}
+
+	return string(leader)
+}
+
+// RemoveExpectedPrimary removes the expected primary key.
+func RemoveExpectedPrimary(client *clientv3.Client, leaderPath string) {
+	// remove expected leader key
+	resp, err := kv.NewSlowLogTxn(client).
+		Then(clientv3.OpDelete(strings.Join([]string{leaderPath, ExpectedPrimary}, "/"))).
+		Commit()
+	if err != nil || !resp.Succeeded {
+		log.Error("change primary error", errs.ZapError(err))
+		return
+	}
+}
+
+// SetExpectedPrimary sets the expected primary key when the current primary has exited.
+func SetExpectedPrimary(client *clientv3.Client, leaderPath string) {
+	leaderRaw, err := etcdutil.GetValue(client, leaderPath)
+	if err != nil {
+		log.Error("[primary] get primary key error", zap.Error(err))
+		return
+	}
+	// write a flag to indicate the current primary has exited
+	resp, err := kv.NewSlowLogTxn(client).
+		Then(
+			clientv3.OpPut(strings.Join([]string{leaderPath, ExpectedPrimary}, "/"), string(leaderRaw)),
+			// indicate the current primary has exited
+			clientv3.OpDelete(leaderPath)).
+		Commit()
+	if err != nil || !resp.Succeeded {
+		log.Error("change primary error", errs.ZapError(err))
+		return
+	}
+}
+
 // PromHandler is a handler to get prometheus metrics.
 func PromHandler() gin.HandlerFunc {
 	return func(c *gin.Context) {