pdms: support primary/transfer api for scheduling and tso

pkg/election/leadership.go

@@ -71,6 +71,9 @@ type Leadership struct {
 	// campaignTimes is used to record the campaign times of the leader within `campaignTimesRecordTimeout`.
 	// It is ordered by time to prevent the leader from campaigning too frequently.
 	campaignTimes []time.Time
+	// primaryWatch is for the primary watch only,
+	// which is used to reuse `Watch` interface in `Leadership`.
+	primaryWatch atomic.Bool
 }
 
 // NewLeadership creates a new Leadership.
@@ -84,17 +87,18 @@ func NewLeadership(client *clientv3.Client, leaderKey, purpose string) *Leadersh
 	return leadership
 }
 
-// getLease gets the lease of leadership, only if leadership is valid,
+// GetLease gets the lease of leadership, only if leadership is valid,
 // i.e. the owner is a true leader, the lease is not nil.
-func (ls *Leadership) getLease() *lease {
+func (ls *Leadership) GetLease() *Lease {
 	l := ls.lease.Load()
 	if l == nil {
 		return nil
 	}
-	return l.(*lease)
+	return l.(*Lease)
 }
 
-func (ls *Leadership) setLease(lease *lease) {
+// SetLease sets the lease of leadership.
+func (ls *Leadership) SetLease(lease *Lease) {
 	ls.lease.Store(lease)
 }
 
@@ -114,6 +118,16 @@ func (ls *Leadership) GetLeaderKey() string {
 	return ls.leaderKey
 }
 
+// SetPrimaryWatch sets the primary watch flag.
+func (ls *Leadership) SetPrimaryWatch(val bool) {
+	ls.primaryWatch.Store(val)
+}
+
+// IsPrimary gets the primary watch flag.
+func (ls *Leadership) IsPrimary() bool {
+	return ls.primaryWatch.Load()
+}
+
 // GetCampaignTimesNum is used to get the campaign times of the leader within `campaignTimesRecordTimeout`.
 // Need to make sure `AddCampaignTimes` is called before this function.
 func (ls *Leadership) GetCampaignTimesNum() int {
@@ -152,18 +166,23 @@ func (ls *Leadership) AddCampaignTimes() {
 func (ls *Leadership) Campaign(leaseTimeout int64, leaderData string, cmps ...clientv3.Cmp) error {
 	ls.leaderValue = leaderData
 	// Create a new lease to campaign
-	newLease := &lease{
+	newLease := &Lease{
 		Purpose: ls.purpose,
 		client:  ls.client,
 		lease:   clientv3.NewLease(ls.client),
 	}
-	ls.setLease(newLease)
+	ls.SetLease(newLease)
 
 	failpoint.Inject("skipGrantLeader", func(val failpoint.Value) {
+		name, ok := val.(string)
+		if name == "" {
+			// return directly when not set the name
+			failpoint.Return(errors.Errorf("failed to grant lease"))
+		}
 		var member pdpb.Member
 		_ = member.Unmarshal([]byte(leaderData))
-		name, ok := val.(string)
 		if ok && member.Name == name {
+			// only return when the name is set and the name is equal to the leader name
 			failpoint.Return(errors.Errorf("failed to grant lease"))
 		}
 	})
@@ -200,12 +219,12 @@ func (ls *Leadership) Keep(ctx context.Context) {
 	ls.keepAliveCancelFuncLock.Lock()
 	ls.keepAliveCtx, ls.keepAliveCancelFunc = context.WithCancel(ctx)
 	ls.keepAliveCancelFuncLock.Unlock()
-	go ls.getLease().KeepAlive(ls.keepAliveCtx)
+	go ls.GetLease().KeepAlive(ls.keepAliveCtx)
 }
 
 // Check returns whether the leadership is still available.
 func (ls *Leadership) Check() bool {
-	return ls != nil && ls.getLease() != nil && !ls.getLease().IsExpired()
+	return ls != nil && ls.GetLease() != nil && !ls.GetLease().IsExpired()
 }
 
 // LeaderTxn returns txn() with a leader comparison to guarantee that
@@ -371,11 +390,19 @@ func (ls *Leadership) Watch(serverCtx context.Context, revision int64) {
 			}
 
 			for _, ev := range wresp.Events {
-				if ev.Type == mvccpb.DELETE {
+				// ensure `{service}/primary/transfer` API will not meet this condition.
+				if ev.Type == mvccpb.DELETE && !ls.IsPrimary() {
 					log.Info("current leadership is deleted",
 						zap.Int64("revision", wresp.Header.Revision), zap.String("leader-key", ls.leaderKey), zap.String("purpose", ls.purpose))
 					return
 				}
+				// ONLY `{service}/primary/transfer` API update primary will meet this condition.
+				if ev.Type == mvccpb.PUT && ls.IsPrimary() {
+					log.Info("current leadership is updated", zap.Int64("revision", wresp.Header.Revision),
+						zap.String("leader-key", ls.leaderKey), zap.ByteString("cur-value", ev.Kv.Value),
+						zap.String("purpose", ls.purpose))
+					return
+				}
 			}
 			revision = wresp.Header.Revision + 1
 		}
@@ -385,13 +412,14 @@ func (ls *Leadership) Watch(serverCtx context.Context, revision int64) {
 
 // Reset does some defer jobs such as closing lease, resetting lease etc.
 func (ls *Leadership) Reset() {
-	if ls == nil || ls.getLease() == nil {
+	if ls == nil || ls.GetLease() == nil {
 		return
 	}
 	ls.keepAliveCancelFuncLock.Lock()
 	if ls.keepAliveCancelFunc != nil {
 		ls.keepAliveCancelFunc()
 	}
 	ls.keepAliveCancelFuncLock.Unlock()
-	ls.getLease().Close()
+	ls.GetLease().Close()
+	ls.SetPrimaryWatch(false)
 }

pkg/election/leadership_test.go

@@ -100,9 +100,9 @@ func TestLeadership(t *testing.T) {
 	leadership2.Keep(ctx)
 
 	// Check the lease.
-	lease1 := leadership1.getLease()
+	lease1 := leadership1.GetLease()
 	re.NotNil(lease1)
-	lease2 := leadership2.getLease()
+	lease2 := leadership2.GetLease()
 	re.NotNil(lease2)
 
 	re.True(lease1.IsExpired())

pkg/election/lease.go

@@ -34,9 +34,9 @@ const (
 	slowRequestTime    = etcdutil.DefaultSlowRequestTime
 )
 
-// lease is used as the low-level mechanism for campaigning and renewing elected leadership.
+// Lease is used as the low-level mechanism for campaigning and renewing elected leadership.
 // The way to gain and maintain leadership is to update and keep the lease alive continuously.
-type lease struct {
+type Lease struct {
 	// purpose is used to show what this election for
 	Purpose string
 	// etcd client and lease
@@ -48,8 +48,17 @@ type lease struct {
 	expireTime   atomic.Value
 }
 
+// NewLease creates a new Lease instance.
+func NewLease(client *clientv3.Client, purpose string) *Lease {
+	return &Lease{
+		Purpose: purpose,
+		client:  client,
+		lease:   clientv3.NewLease(client),
+	}
+}
+
 // Grant uses `lease.Grant` to initialize the lease and expireTime.
-func (l *lease) Grant(leaseTimeout int64) error {
+func (l *Lease) Grant(leaseTimeout int64) error {
 	if l == nil {
 		return errs.ErrEtcdGrantLease.GenWithStackByCause("lease is nil")
 	}
@@ -71,7 +80,7 @@ func (l *lease) Grant(leaseTimeout int64) error {
 }
 
 // Close releases the lease.
-func (l *lease) Close() error {
+func (l *Lease) Close() error {
 	if l == nil {
 		return nil
 	}
@@ -92,15 +101,15 @@ func (l *lease) Close() error {
 
 // IsExpired checks if the lease is expired. If it returns true,
 // current leader should step down and try to re-elect again.
-func (l *lease) IsExpired() bool {
+func (l *Lease) IsExpired() bool {
 	if l == nil || l.expireTime.Load() == nil {
 		return true
 	}
 	return time.Now().After(l.expireTime.Load().(time.Time))
 }
 
 // KeepAlive auto renews the lease and update expireTime.
-func (l *lease) KeepAlive(ctx context.Context) {
+func (l *Lease) KeepAlive(ctx context.Context) {
 	defer logutil.LogPanic()
 
 	if l == nil {
@@ -109,6 +118,7 @@ func (l *lease) KeepAlive(ctx context.Context) {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 	timeCh := l.keepAliveWorker(ctx, l.leaseTimeout/3)
+	defer log.Info("lease keep alive stopped", zap.String("purpose", l.Purpose))
 
 	var maxExpire time.Time
 	timer := time.NewTimer(l.leaseTimeout)
@@ -146,7 +156,7 @@ func (l *lease) KeepAlive(ctx context.Context) {
 }
 
 // Periodically call `lease.KeepAliveOnce` and post back latest received expire time into the channel.
-func (l *lease) keepAliveWorker(ctx context.Context, interval time.Duration) <-chan time.Time {
+func (l *Lease) keepAliveWorker(ctx context.Context, interval time.Duration) <-chan time.Time {
 	ch := make(chan time.Time)
 
 	go func() {

pkg/election/lease_test.go

@@ -30,12 +30,12 @@ func TestLease(t *testing.T) {
 	defer clean()
 
 	// Create the lease.
-	lease1 := &lease{
+	lease1 := &Lease{
 		Purpose: "test_lease_1",
 		client:  client,
 		lease:   clientv3.NewLease(client),
 	}
-	lease2 := &lease{
+	lease2 := &Lease{
 		Purpose: "test_lease_2",
 		client:  client,
 		lease:   clientv3.NewLease(client),
@@ -95,7 +95,7 @@ func TestLeaseKeepAlive(t *testing.T) {
 	defer clean()
 
 	// Create the lease.
-	lease := &lease{
+	lease := &Lease{
 		Purpose: "test_lease",
 		client:  client,
 		lease:   clientv3.NewLease(client),

pkg/mcs/discovery/discover.go

@@ -15,12 +15,16 @@
 package discovery
 
 import (
+	"math/rand"
 	"strconv"
+	"time"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/log"
+	"github.com/tikv/pd/pkg/election"
 	"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"
@@ -77,3 +81,66 @@ func GetMSMembers(serviceName string, client *clientv3.Client) ([]ServiceRegistr
 
 	return nil, errors.Errorf("unknown service name %s", serviceName)
 }
+
+// TransferPrimary transfers the primary of the specified service.
+// keyspaceGroupID is optional, only used for TSO service.
+func TransferPrimary(client *clientv3.Client, lease *election.Lease, serviceName,
+	oldPrimaryAddr, newPrimary string, keyspaceGroupID uint32) error {
+	if lease == nil {
+		return errors.New("current lease is nil, please check leadership")
+	}
+	log.Info("try to transfer primary", zap.String("service", serviceName), zap.String("from", oldPrimaryAddr), 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 {
+		return errors.Errorf("no valid secondary to transfer primary, the only member is %s", entries[0].Name)
+	}
+
+	var primaryIDs []string
+	for _, member := range entries {
+		// TODO: judged by `addr` and `name` now, should unify them to `name` in the future.
+		if (newPrimary == "" && member.ServiceAddr != oldPrimaryAddr) || (newPrimary != "" && member.Name == newPrimary) {
+			primaryIDs = append(primaryIDs, member.ServiceAddr)
+		}
+	}
+	if len(primaryIDs) == 0 {
+		return errors.Errorf("no valid secondary to transfer primary, from %s to %s", oldPrimaryAddr, newPrimary)
+	}
+
+	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)
+	}
+
+	// update expected primary flag
+	grantResp, err := client.Grant(client.Ctx(), utils.DefaultLeaderLease)
+	if err != nil {
+		return errors.Errorf("failed to grant lease for expected primary, err: %v", err)
+	}
+
+	// revoke current primary's lease to ensure keepalive goroutine of primary exits.
+	if err := lease.Close(); err != nil {
+		return errors.Errorf("failed to revoke current primary's lease: %v", err)
+	}
+
+	var primaryPath string
+	switch serviceName {
+	case utils.SchedulingServiceName:
+		primaryPath = endpoint.SchedulingPrimaryPath(clusterID)
+	case utils.TSOServiceName:
+		tsoRootPath := endpoint.TSOSvcRootPath(clusterID)
+		primaryPath = endpoint.KeyspaceGroupPrimaryPath(tsoRootPath, keyspaceGroupID)
+	}
+	_, err = utils.MarkExpectedPrimaryFlag(client, primaryPath, primaryIDs[nextPrimaryID], grantResp.ID)
+	if err != nil {
+		return errors.Errorf("failed to mark expected primary flag for %s, err: %v", serviceName, err)
+	}
+	return nil
+}

pkg/mcs/scheduling/server/apis/v1/api.go

@@ -30,6 +30,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/log"
 	"github.com/tikv/pd/pkg/errs"
+	"github.com/tikv/pd/pkg/mcs/discovery"
 	scheserver "github.com/tikv/pd/pkg/mcs/scheduling/server"
 	mcsutils "github.com/tikv/pd/pkg/mcs/utils"
 	"github.com/tikv/pd/pkg/response"
@@ -120,6 +121,7 @@ func NewService(srv *scheserver.Service) *Service {
 	s.RegisterHotspotRouter()
 	s.RegisterRegionsRouter()
 	s.RegisterStoresRouter()
+	s.RegisterPrimaryRouter()
 	return s
 }
 
@@ -226,6 +228,12 @@ func (s *Service) RegisterConfigRouter() {
 	regions.GET("/:id/labels", getRegionLabels)
 }
 
+// RegisterPrimaryRouter registers the router of the config handler.
+func (s *Service) RegisterPrimaryRouter() {
+	router := s.root.Group("primary")
+	router.POST("transfer", transferPrimary)
+}
+
 // @Tags     admin
 // @Summary  Change the log level.
 // @Produce  json
@@ -1478,3 +1486,32 @@ func getRegionByID(c *gin.Context) {
 	}
 	c.Data(http.StatusOK, "application/json", b)
 }
+
+// TransferPrimary transfers the primary member.
+// @Tags     primary
+// @Summary  Transfer the primary member of the specified service.
+// @Produce  json
+// @Param    service     path    string  true  "service name"
+// @Param    new_primary body   string  false "new primary name"
+// @Success  200  string  string
+// @Router   /primary/transfer [post]
+func transferPrimary(c *gin.Context) {
+	svr := c.MustGet(multiservicesapi.ServiceContextKey).(*scheserver.Server)
+	var input map[string]string
+	if err := c.BindJSON(&input); err != nil {
+		c.String(http.StatusBadRequest, err.Error())
+		return
+	}
+
+	newPrimary := ""
+	if v, ok := input["new_primary"]; ok {
+		newPrimary = v
+	}
+
+	if err := discovery.TransferPrimary(svr.GetClient(), svr.GetParticipant().GetExpectedPrimaryLease(),
+		mcsutils.SchedulingServiceName, svr.GetAddr(), newPrimary, 0); err != nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
+		return
+	}
+	c.IndentedJSON(http.StatusOK, "success")
+}