tso: add basic election logic for Local TSO Allocator

server/server.go

@@ -349,7 +349,7 @@ func (s *Server) startServer(ctx context.Context) error {
 	s.member.SetMemberGitHash(s.member.ID(), versioninfo.PDGitHash)
 	s.idAllocator = id.NewAllocatorImpl(s.client, s.rootPath, s.member.MemberValue())
 	s.tsoAllocatorManager = tso.NewAllocatorManager(
-		s.member.Etcd(), s.client, s.rootPath, s.cfg.TsoSaveInterval.Duration,
+		s.member, s.rootPath, s.cfg.TsoSaveInterval.Duration,
 		func() time.Duration { return s.persistOptions.GetMaxResetTSGap() },
 	)
 	kvBase := kv.NewEtcdKVBase(s.client, s.rootPath)
@@ -688,6 +688,11 @@ func (s *Server) GetAllocator() *id.AllocatorImpl {
 	return s.idAllocator
 }
 
+// GetTSOAllocatorManager returns the manager of TSO Allocator.
+func (s *Server) GetTSOAllocatorManager() *tso.AllocatorManager {
+	return s.tsoAllocatorManager
+}
+
 // Name returns the unique etcd Name for this server in etcd cluster.
 func (s *Server) Name() string {
 	return s.cfg.Name
@@ -941,6 +946,11 @@ func (s *Server) GetSecurityConfig() *grpcutil.SecurityConfig {
 	return &s.cfg.Security
 }
 
+// GetServerRootPath returns the server root path.
+func (s *Server) GetServerRootPath() string {
+	return s.rootPath
+}
+
 // GetClusterRootPath returns the cluster root path.
 func (s *Server) GetClusterRootPath() string {
 	return path.Join(s.rootPath, "raft")

server/tso/allocator_manager.go

@@ -24,13 +24,16 @@ import (
 	"github.com/pingcap/log"
 	"github.com/tikv/pd/pkg/errs"
 	"github.com/tikv/pd/server/election"
-	"go.etcd.io/etcd/clientv3"
-	"go.etcd.io/etcd/embed"
+	"github.com/tikv/pd/server/member"
 	"go.uber.org/zap"
 )
 
-// GlobalDCLocation is the Global TSO Allocator's dc-location label.
-const GlobalDCLocation = "global"
+const (
+	// GlobalDCLocation is the Global TSO Allocator's dc-location label.
+	GlobalDCLocation            = "global"
+	leaderTickInterval          = 50 * time.Millisecond
+	defaultAllocatorLeaderLease = 3
+)
 
 type allocatorGroup struct {
 	dcLocation string
@@ -63,21 +66,19 @@ type AllocatorManager struct {
 	//   2. Local TSO Allocator, servers for DC-level transactions.
 	// dc-location/global (string) -> TSO Allocator
 	allocatorGroups map[string]*allocatorGroup
-	// etcd and its client
-	etcd   *embed.Etcd
-	client *clientv3.Client
+	// for election use
+	member *member.Member
 	// tso config
 	rootPath      string
 	saveInterval  time.Duration
 	maxResetTSGap func() time.Duration
 }
 
 // NewAllocatorManager creates a new TSO Allocator Manager.
-func NewAllocatorManager(etcd *embed.Etcd, client *clientv3.Client, rootPath string, saveInterval time.Duration, maxResetTSGap func() time.Duration) *AllocatorManager {
+func NewAllocatorManager(m *member.Member, rootPath string, saveInterval time.Duration, maxResetTSGap func() time.Duration) *AllocatorManager {
 	allocatorManager := &AllocatorManager{
 		allocatorGroups: make(map[string]*allocatorGroup),
-		etcd:            etcd,
-		client:          client,
+		member:          m,
 		rootPath:        rootPath,
 		saveInterval:    saveInterval,
 		maxResetTSGap:   maxResetTSGap,
@@ -95,27 +96,125 @@ func (am *AllocatorManager) getAllocatorPath(dcLocation string) string {
 
 // SetUpAllocator is used to set up an allocator, which will initialize the allocator and put it into allocator daemon.
 func (am *AllocatorManager) SetUpAllocator(parentCtx context.Context, parentCancel context.CancelFunc, dcLocation string, leadership *election.Leadership) error {
+	var allocator Allocator
+	if dcLocation == GlobalDCLocation {
+		allocator = NewGlobalTSOAllocator(leadership, am.getAllocatorPath(dcLocation), am.saveInterval, am.maxResetTSGap)
+	} else {
+		allocator = NewLocalTSOAllocator(am.member, leadership, am.getAllocatorPath(dcLocation), dcLocation, am.saveInterval, am.maxResetTSGap)
+	}
 	am.Lock()
 	defer am.Unlock()
+	// Update or create a new allocatorGroup
+	am.allocatorGroups[dcLocation] = &allocatorGroup{
+		dcLocation:   dcLocation,
+		parentCtx:    parentCtx,
+		parentCancel: parentCancel,
+		leadership:   leadership,
+		allocator:    allocator,
+	}
+	// Different kinds of allocators have different setup works to do
 	switch dcLocation {
+	// For Global TSO Allocator
 	case GlobalDCLocation:
-		am.allocatorGroups[dcLocation] = &allocatorGroup{
-			dcLocation:   dcLocation,
-			parentCtx:    parentCtx,
-			parentCancel: parentCancel,
-			leadership:   leadership,
-			allocator:    NewGlobalTSOAllocator(leadership, am.getAllocatorPath(dcLocation), am.saveInterval, am.maxResetTSGap),
-		}
+		// Because Global TSO Allocator only depends on PD leader's leadership,
+		// so we can directly initialize it here.
 		if err := am.allocatorGroups[dcLocation].allocator.Initialize(); err != nil {
 			return err
 		}
 		am.allocatorGroups[dcLocation].isInitialized = true
+	// For Local TSO Allocator
 	default:
-		// Todo: set up a Local TSO Allocator
+		// Join in a Local TSO Allocator election
+		localTSOAllocator, _ := allocator.(*LocalTSOAllocator)
+		go am.allocatorLeaderLoop(parentCtx, localTSOAllocator)
 	}
 	return nil
 }
 
+// similar logic with leaderLoop in server/server.go
+func (am *AllocatorManager) allocatorLeaderLoop(parentCtx context.Context, allocator *LocalTSOAllocator) {
+	for {
+		select {
+		case <-parentCtx.Done():
+			log.Info("server is closed, return local tso allocator leader loop",
+				zap.String("dc-location", allocator.dcLocation),
+				zap.String("local-tso-allocator-name", am.member.Member().Name))
+			return
+		default:
+		}
+
+		allocatorLeader, rev, checkAgain := allocator.CheckAllocatorLeader()
+		if checkAgain {
+			continue
+		}
+		if allocatorLeader != nil {
+			log.Info("start to watch allocator leader",
+				zap.Stringer(fmt.Sprintf("%s-allocator-leader", allocator.dcLocation), allocatorLeader),
+				zap.String("local-tso-allocator-name", am.member.Member().Name))
+			// WatchAllocatorLeader will keep looping and never return unless the Local TSO Allocator leader has changed.
+			allocator.WatchAllocatorLeader(parentCtx, allocatorLeader, rev)
+			log.Info("local tso allocator leader has changed, try to re-campaign a local tso allocator leader",
+				zap.String("dc-location", allocator.dcLocation))
+		}
+		am.campaignAllocatorLeader(parentCtx, allocator)
+	}
+}
+
+func (am *AllocatorManager) campaignAllocatorLeader(parentCtx context.Context, allocator *LocalTSOAllocator) {
+	log.Info("start to campaign local tso allocator leader",
+		zap.String("dc-location", allocator.dcLocation),
+		zap.String("campaign-local-tso-allocator-leader-name", am.member.Member().Name))
+	if err := allocator.CampaignAllocatorLeader(defaultAllocatorLeaderLease); err != nil {
+		log.Error("failed to campaign local tso allocator leader", errs.ZapError(err))
+		return
+	}
+
+	// Start keepalive the Local TSO Allocator leadership and enable Local TSO service.
+	ctx, cancel := context.WithCancel(parentCtx)
+	defer cancel()
+	// maintain the Local TSO Allocator leader
+	go allocator.KeepAllocatorLeader(ctx)
+	log.Info("campaign local tso allocator leader ok",
+		zap.String("dc-location", allocator.dcLocation),
+		zap.String("campaign-local-tso-allocator-leader-name", am.member.Member().Name))
+
+	log.Info("initialize the local TSO allocator", zap.String("dc-location", allocator.dcLocation))
+	if err := allocator.Initialize(); err != nil {
+		log.Error("failed to initialize the local TSO allocator", errs.ZapError(err))
+		return
+	}
+	am.Lock()
+	am.allocatorGroups[allocator.dcLocation].isInitialized = true
+	am.Unlock()
+	allocator.EnableAllocatorLeader()
+	log.Info("local tso allocator leader is ready to serve",
+		zap.String("dc-location", allocator.dcLocation),
+		zap.String("campaign-local-tso-allocator-leader-name", am.member.Member().Name))
+
+	leaderTicker := time.NewTicker(leaderTickInterval)
+	defer leaderTicker.Stop()
+
+	for {
+		select {
+		case <-leaderTicker.C:
+			if !allocator.IsStillAllocatorLeader() {
+				log.Info("no longer a local tso allocator leader because lease has expired, local tso allocator leader will step down",
+					zap.String("dc-location", allocator.dcLocation))
+				return
+			}
+		case <-ctx.Done():
+			// Server is closed and it should return nil.
+			log.Info("server is closed, reset the local tso allocator", zap.String("dc-location", allocator.dcLocation))
+			am.Lock()
+			allocator.Reset()
+			am.allocatorGroups[allocator.dcLocation].leadership.Reset()
+			am.allocatorGroups[allocator.dcLocation].isInitialized = false
+			am.Unlock()
+			return
+		}
+	}
+}
+
 // GetAllocator get the allocator by dc-location.
 func (am *AllocatorManager) GetAllocator(dcLocation string) (Allocator, error) {
 	am.RLock()
@@ -127,6 +226,26 @@ func (am *AllocatorManager) GetAllocator(dcLocation string) (Allocator, error) {
 	return allocatorGroup.allocator, nil
 }
 
+// GetAllocators get all allocators with some filters.
+func (am *AllocatorManager) GetAllocators(withGlobal, withLocal, withInitialized bool) []Allocator {
+	am.RLock()
+	defer am.RUnlock()
+	var allocators []Allocator
+	for dcLocation, allocatorGroup := range am.allocatorGroups {
+		if !withGlobal && dcLocation == GlobalDCLocation {
+			continue
+		}
+		if !withLocal && dcLocation != GlobalDCLocation {
+			continue
+		}
+		if withInitialized && !allocatorGroup.isInitialized {
+			continue
+		}
+		allocators = append(allocators, allocatorGroup.allocator)
+	}
+	return allocators
+}
+
 func (am *AllocatorManager) getAllocatorGroups() []*allocatorGroup {
 	am.RLock()
 	defer am.RUnlock()
@@ -149,8 +268,11 @@ func (am *AllocatorManager) AllocatorDaemon(serverCtx context.Context) {
 			allocatorGroups := am.getAllocatorGroups()
 			// Update each allocator concurrently
 			for _, ag := range allocatorGroups {
+				am.RLock()
 				// Filter allocators without leadership and uninitialized
-				if ag.isInitialized && ag.leadership.Check() {
+				notFilterd := ag.isInitialized && ag.leadership.Check()
+				am.RUnlock()
+				if notFilterd {
 					am.wg.Add(1)
 					go am.updateAllocator(ag)
 				}
@@ -180,7 +302,7 @@ func (am *AllocatorManager) updateAllocator(ag *allocatorGroup) {
 		return
 	}
 	if err := ag.allocator.UpdateTSO(); err != nil {
-		log.Warn("failed to update allocator's timestamp", zap.String("dc-location", ag.dcLocation), zap.Error(err))
+		log.Warn("failed to update allocator's timestamp", zap.String("dc-location", ag.dcLocation), errs.ZapError(err))
 		ag.parentCancel()
 		return
 	}