*: refactor dynamic tso switch

pkg/keyspace/keyspace_test.go

@@ -47,9 +47,11 @@ func TestMain(m *testing.M) {
 }
 
 const (
-	testConfig  = "test config"
-	testConfig1 = "config_entry_1"
-	testConfig2 = "config_entry_2"
+	testConfig   = "test config"
+	testConfig1  = "config_entry_1"
+	testConfig2  = "config_entry_2"
+	testGroupID  = "tso_keyspace_group_id"
+	testUserKind = "user_kind"
 )
 
 type keyspaceTestSuite struct {
@@ -128,8 +130,10 @@ func makeCreateKeyspaceRequests(count int) []*CreateKeyspaceRequest {
 		requests[i] = &CreateKeyspaceRequest{
 			Name: fmt.Sprintf("test_keyspace_%d", i),
 			Config: map[string]string{
-				testConfig1: "100",
-				testConfig2: "200",
+				testConfig1:  "100",
+				testConfig2:  "200",
+				testGroupID:  "0",
+				testUserKind: "basic",
 			},
 			CreateTime: now,
 		}

pkg/mcs/discovery/register.go

@@ -63,6 +63,7 @@ func (sr *ServiceRegister) Register() error {
 		sr.cancel()
 		return fmt.Errorf("put the key with lease %s failed: %v", sr.key, err)
 	}
+	log.Info("register service", zap.String("key", sr.key))
 	kresp, err := sr.cli.KeepAlive(sr.ctx, id)
 	if err != nil {
 		sr.cancel()

pkg/mcs/registry/registry.go

@@ -80,18 +80,18 @@ func (r *ServiceRegistry) InstallAllRESTHandler(srv bs.Server, h map[string]http
 		serviceName := createServiceName(prefix, name)
 		if l, ok := r.services[serviceName]; ok {
 			if err := l.RegisterRESTHandler(h); err != nil {
-				log.Error("register restful PD failed", zap.String("prefix", prefix), zap.String("service-name", name), zap.Error(err))
+				log.Error("register restful service failed", zap.String("prefix", prefix), zap.String("service-name", name), zap.Error(err))
 			} else {
-				log.Info("restful PD already registered", zap.String("prefix", prefix), zap.String("service-name", name))
+				log.Info("restful service already registered", zap.String("prefix", prefix), zap.String("service-name", name))
 			}
 			continue
 		}
 		l := builder(srv)
 		r.services[serviceName] = l
 		if err := l.RegisterRESTHandler(h); err != nil {
-			log.Error("register restful PD failed", zap.String("prefix", prefix), zap.String("service-name", name), zap.Error(err))
+			log.Error("register restful service failed", zap.String("prefix", prefix), zap.String("service-name", name), zap.Error(err))
 		} else {
-			log.Info("restful PD registered successfully", zap.String("prefix", prefix), zap.String("service-name", name))
+			log.Info("restful service registered successfully", zap.String("prefix", prefix), zap.String("service-name", name))
 		}
 	}
 }

pkg/storage/endpoint/tso.go

@@ -21,10 +21,13 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/log"
 
 	"github.com/tikv/pd/pkg/election"
 	"github.com/tikv/pd/pkg/errs"
+	keyspaceconstant "github.com/tikv/pd/pkg/keyspace/constant"
+	mcsconstant "github.com/tikv/pd/pkg/mcs/utils/constant"
 	"github.com/tikv/pd/pkg/storage/kv"
 	"github.com/tikv/pd/pkg/utils/keypath"
 	"github.com/tikv/pd/pkg/utils/typeutil"
@@ -33,12 +36,16 @@ import (
 // TSOStorage is the interface for timestamp storage.
 type TSOStorage interface {
 	LoadTimestamp(groupID uint32) (time.Time, error)
-	SaveTimestamp(ctx context.Context, groupID uint32, ts time.Time, leadership *election.Leadership) error
+	SaveTimestamp(ctx context.Context, groupID uint32, ts time.Time, leadership *election.Leadership, checkTSOPrimary bool) error
 	DeleteTimestamp(ctx context.Context, groupID uint32) error
 }
 
 var _ TSOStorage = (*StorageEndpoint)(nil)
 
+// ErrTSOServicePrimaryExists indicates that embedded PD TSO must stop updating
+// its timestamp window because a TSO service primary has been elected.
+var ErrTSOServicePrimaryExists = errors.New("tso microservice primary exists, pd must yield embedded tso")
+
 // LoadTimestamp retrieves the last saved TSO timestamp from etcd.
 // Before switching back from the TSO microservice to the PD leader,
 // we must ensure that all keyspace groups are merged into the default
@@ -68,20 +75,28 @@ func (se *StorageEndpoint) LoadTimestamp(groupID uint32) (time.Time, error) {
 
 // SaveTimestamp saves the timestamp to the storage. The leadership is used to check if the current server is leader
 // before saving the timestamp to ensure a strong consistency for persistence of the TSO timestamp window.
-func (se *StorageEndpoint) SaveTimestamp(ctx context.Context, groupID uint32, ts time.Time, leadership *election.Leadership) error {
+func (se *StorageEndpoint) SaveTimestamp(ctx context.Context, groupID uint32, ts time.Time, leadership *election.Leadership, checkTSOPrimary bool) error {
 	logFilds := []zap.Field{
 		zap.Uint32("group-id", groupID),
 		zap.Time("ts", ts),
 		zap.String("leader-key", leadership.GetLeaderKey()),
 		zap.String("expected-leader-value", leadership.GetLeaderValue()),
+		zap.Bool("check-tso-primary", checkTSOPrimary),
 	}
 	log.Debug("saving timestamp to the storage", logFilds...)
 	// The PD leadership or TSO primary will always be granted first before the TSO timestamp window is saved.
 	// So we here check whether the leader value is filled to see if the requirement is met.
 	if len(leadership.GetLeaderValue()) == 0 {
 		return errors.Errorf("%s due to leadership has not been granted yet", errs.NotLeaderErr)
 	}
-	return se.RunInTxn(ctx, func(txn kv.Txn) error {
+	var tsoPrimaryPath string
+	if checkTSOPrimary {
+		tsoPrimaryPath = keypath.ElectionPath(&keypath.MsParam{
+			ServiceName: mcsconstant.TSOServiceName,
+			GroupID:     keyspaceconstant.DefaultKeyspaceGroupID,
+		})
+	}
+	err := se.RunInTxn(ctx, func(txn kv.Txn) error {
 		// Ensure the current server is leader by reading and comparing the leader value.
 		leaderValue, err := txn.Load(leadership.GetLeaderKey())
 		if err != nil {
@@ -92,6 +107,16 @@ func (se *StorageEndpoint) SaveTimestamp(ctx context.Context, groupID uint32, ts
 			return errors.Errorf("%s due to leader value does not match, current: %s, expected: %s", errs.NotLeaderErr, leaderValue, expected)
 		}
 
+		if checkTSOPrimary {
+			tsoPrimary, err := txn.Load(tsoPrimaryPath)
+			if err != nil {
+				return err
+			}
+			if len(tsoPrimary) != 0 {
+				return ErrTSOServicePrimaryExists
+			}
+		}
+
 		value, err := txn.Load(keypath.TimestampPath(groupID))
 		if err != nil {
 			return err
@@ -108,9 +133,17 @@ func (se *StorageEndpoint) SaveTimestamp(ctx context.Context, groupID uint32, ts
 		if previousTS != typeutil.ZeroTime && typeutil.SubRealTimeByWallClock(ts, previousTS) <= 0 {
 			return errors.Errorf("saving timestamp %d is less than or equal to the previous one %d", ts.UnixNano(), previousTS.UnixNano())
 		}
+		failpoint.InjectCall("beforeSaveTimestampTxnCommit")
 		data := typeutil.Uint64ToBytes(uint64(ts.UnixNano()))
 		return txn.Save(keypath.TimestampPath(groupID), string(data))
 	})
+	if checkTSOPrimary && errors.ErrorEqual(err, errs.ErrEtcdTxnConflict) {
+		tsoPrimary, loadErr := se.Load(tsoPrimaryPath)
+		if loadErr == nil && len(tsoPrimary) != 0 {
+			return ErrTSOServicePrimaryExists
+		}
+	}
+	return err
 }
 
 // DeleteTimestamp deletes the timestamp from the storage.

pkg/storage/storage_tso_test.go

@@ -25,7 +25,11 @@ import (
 
 	"github.com/tikv/pd/pkg/election"
 	"github.com/tikv/pd/pkg/errs"
+	keyspaceconstant "github.com/tikv/pd/pkg/keyspace/constant"
+	mcsconstant "github.com/tikv/pd/pkg/mcs/utils/constant"
+	storageendpoint "github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/utils/etcdutil"
+	"github.com/tikv/pd/pkg/utils/keypath"
 )
 
 const (
@@ -55,7 +59,7 @@ func TestSaveTimestampWithTimeout(t *testing.T) {
 	defer clean()
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
-	err := storage.SaveTimestamp(ctx, testGroupID, time.Now().Round(0), leadership)
+	err := storage.SaveTimestamp(ctx, testGroupID, time.Now().Round(0), leadership, false)
 	re.ErrorIs(err, context.DeadlineExceeded)
 }
 
@@ -64,7 +68,7 @@ func TestSaveLoadTimestamp(t *testing.T) {
 	storage, clean, leadership := prepare(t)
 	defer clean()
 	expectedTS := time.Now().Round(0)
-	err := storage.SaveTimestamp(defaultContext, testGroupID, expectedTS, leadership)
+	err := storage.SaveTimestamp(defaultContext, testGroupID, expectedTS, leadership, false)
 	re.NoError(err)
 	ts, err := storage.LoadTimestamp(testGroupID)
 	re.NoError(err)
@@ -76,11 +80,11 @@ func TestTimestampTxn(t *testing.T) {
 	storage, clean, leadership := prepare(t)
 	defer clean()
 	globalTS1 := time.Now().Round(0)
-	err := storage.SaveTimestamp(defaultContext, testGroupID, globalTS1, leadership)
+	err := storage.SaveTimestamp(defaultContext, testGroupID, globalTS1, leadership, false)
 	re.NoError(err)
 
 	globalTS2 := globalTS1.Add(-time.Millisecond).Round(0)
-	err = storage.SaveTimestamp(defaultContext, testGroupID, globalTS2, leadership)
+	err = storage.SaveTimestamp(defaultContext, testGroupID, globalTS2, leadership, false)
 	re.Error(err)
 
 	ts, err := storage.LoadTimestamp(testGroupID)
@@ -95,13 +99,13 @@ func TestSaveTimestampWithLeaderCheck(t *testing.T) {
 
 	// testLeaderKey -> testLeaderValue
 	globalTS := time.Now().Round(0)
-	err := storage.SaveTimestamp(defaultContext, testGroupID, globalTS, leadership)
+	err := storage.SaveTimestamp(defaultContext, testGroupID, globalTS, leadership, false)
 	re.NoError(err)
 	ts, err := storage.LoadTimestamp(testGroupID)
 	re.NoError(err)
 	re.Equal(globalTS, ts)
 
-	err = storage.SaveTimestamp(context.Background(), testGroupID, globalTS.Add(time.Second), &election.Leadership{})
+	err = storage.SaveTimestamp(context.Background(), testGroupID, globalTS.Add(time.Second), &election.Leadership{}, false)
 	re.True(errs.IsLeaderChanged(err))
 	ts, err = storage.LoadTimestamp(testGroupID)
 	re.NoError(err)
@@ -110,7 +114,7 @@ func TestSaveTimestampWithLeaderCheck(t *testing.T) {
 	// testLeaderKey -> ""
 	err = storage.Save(leadership.GetLeaderKey(), "")
 	re.NoError(err)
-	err = storage.SaveTimestamp(defaultContext, testGroupID, globalTS.Add(time.Second), leadership)
+	err = storage.SaveTimestamp(defaultContext, testGroupID, globalTS.Add(time.Second), leadership, false)
 	re.True(errs.IsLeaderChanged(err))
 	ts, err = storage.LoadTimestamp(testGroupID)
 	re.NoError(err)
@@ -119,7 +123,7 @@ func TestSaveTimestampWithLeaderCheck(t *testing.T) {
 	// testLeaderKey -> non-existent
 	err = storage.Remove(leadership.GetLeaderKey())
 	re.NoError(err)
-	err = storage.SaveTimestamp(defaultContext, testGroupID, globalTS.Add(time.Second), leadership)
+	err = storage.SaveTimestamp(defaultContext, testGroupID, globalTS.Add(time.Second), leadership, false)
 	re.True(errs.IsLeaderChanged(err))
 	ts, err = storage.LoadTimestamp(testGroupID)
 	re.NoError(err)
@@ -129,9 +133,77 @@ func TestSaveTimestampWithLeaderCheck(t *testing.T) {
 	err = storage.Save(leadership.GetLeaderKey(), testLeaderValue)
 	re.NoError(err)
 	globalTS = globalTS.Add(time.Second)
-	err = storage.SaveTimestamp(defaultContext, testGroupID, globalTS, leadership)
+	err = storage.SaveTimestamp(defaultContext, testGroupID, globalTS, leadership, false)
 	re.NoError(err)
 	ts, err = storage.LoadTimestamp(testGroupID)
 	re.NoError(err)
 	re.Equal(globalTS, ts)
 }
+
+func TestSaveTimestampCheckTSOPrimary(t *testing.T) {
+	re := require.New(t)
+	storage, clean, leadership := prepare(t)
+	defer clean()
+
+	globalTS := time.Now().Round(0)
+	err := storage.SaveTimestamp(defaultContext, testGroupID, globalTS, leadership, true)
+	re.NoError(err)
+
+	tsoPrimaryPath := keypath.ElectionPath(&keypath.MsParam{
+		ServiceName: mcsconstant.TSOServiceName,
+		GroupID:     keyspaceconstant.DefaultKeyspaceGroupID,
+	})
+	err = storage.Save(tsoPrimaryPath, "tso-primary")
+	re.NoError(err)
+
+	err = storage.SaveTimestamp(defaultContext, testGroupID, globalTS.Add(time.Second), leadership, true)
+	re.ErrorIs(err, storageendpoint.ErrTSOServicePrimaryExists)
+
+	ts, err := storage.LoadTimestamp(testGroupID)
+	re.NoError(err)
+	re.Equal(globalTS, ts)
+
+	err = storage.SaveTimestamp(defaultContext, testGroupID, globalTS.Add(time.Second), leadership, false)
+	re.NoError(err)
+}
+
+func TestSaveTimestampFencedByConcurrentTSOPrimary(t *testing.T) {
+	re := require.New(t)
+	storage, clean, leadership := prepare(t)
+	defer clean()
+
+	globalTS := time.Now().Round(0)
+	re.NoError(storage.SaveTimestamp(defaultContext, testGroupID, globalTS, leadership, true))
+
+	entered := make(chan struct{})
+	resume := make(chan struct{})
+	re.NoError(failpoint.EnableCall("github.com/tikv/pd/pkg/storage/endpoint/beforeSaveTimestampTxnCommit", func() {
+		close(entered)
+		<-resume
+	}))
+	defer func() {
+		re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/storage/endpoint/beforeSaveTimestampTxnCommit"))
+	}()
+
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- storage.SaveTimestamp(defaultContext, testGroupID, globalTS.Add(time.Second), leadership, true)
+	}()
+	select {
+	case <-entered:
+	case <-time.After(time.Second):
+		t.Fatal("timestamp transaction did not reach the commit fence")
+	}
+
+	tsoPrimaryPath := keypath.ElectionPath(&keypath.MsParam{
+		ServiceName: mcsconstant.TSOServiceName,
+		GroupID:     keyspaceconstant.DefaultKeyspaceGroupID,
+	})
+	re.NoError(storage.Save(tsoPrimaryPath, "tso-primary"))
+	close(resume)
+	re.ErrorIs(<-errCh, storageendpoint.ErrTSOServicePrimaryExists)
+
+	ts, err := storage.LoadTimestamp(testGroupID)
+	re.NoError(err)
+	re.Equal(globalTS, ts)
+}

pkg/tso/allocator.go

@@ -78,6 +78,7 @@ func NewAllocator(
 	member member.Election,
 	storage endpoint.TSOStorage,
 	cfg Config,
+	checkTSOPrimary func() bool,
 ) *Allocator {
 	ctx, cancel := context.WithCancel(ctx)
 	keyspaceGroupIDStr := strconv.FormatUint(uint64(keyspaceGroupID), 10)
@@ -95,6 +96,7 @@ func NewAllocator(
 			updatePhysicalInterval: cfg.GetTSOUpdatePhysicalInterval(),
 			maxResetTSGap:          cfg.GetMaxResetTSGap,
 			tsoMux:                 &tsoObject{},
+			checkTSOPrimary:        checkTSOPrimary,
 			metrics:                newTSOMetrics(keyspaceGroupIDStr),
 		},
 		tsoAllocatorRoleGauge: tsoAllocatorRole.WithLabelValues(keyspaceGroupIDStr),
@@ -131,6 +133,11 @@ func (a *Allocator) allocatorUpdater() {
 				continue
 			}
 			if err := a.UpdateTSO(); err != nil {
+				if errors.Is(err, endpoint.ErrTSOServicePrimaryExists) {
+					log.Info("stop embedded TSO allocator because TSO service primary exists", a.logFields...)
+					a.Reset(false)
+					continue
+				}
 				log.Warn("failed to update allocator's timestamp, resetting the TSO allocator with leadership resignation",
 					append(a.logFields, errs.ZapError(err))...)
 				a.Reset(true)
@@ -173,6 +180,9 @@ func (a *Allocator) UpdateTSO() (err error) {
 		if err == nil {
 			return nil
 		}
+		if errors.Is(err, endpoint.ErrTSOServicePrimaryExists) {
+			return err
+		}
 		log.Warn("try to update the tso but failed",
 			zap.Int("retry-count", i),
 			zap.Duration("retry-interval", updateTSORetryInterval),

pkg/tso/keyspace_group_manager.go

@@ -804,7 +804,7 @@ func (kgm *KeyspaceGroupManager) updateKeyspaceGroup(group *endpoint.KeyspaceGro
 		})
 	}
 	// Initialize all kinds of maps.
-	allocator := NewAllocator(kgm.ctx, group.ID, participant, kgm.storage, kgm.cfg)
+	allocator := NewAllocator(kgm.ctx, group.ID, participant, kgm.storage, kgm.cfg, nil)
 	allocator.startPrimaryElectionLoop()
 	log.Info("created tso allocator",
 		zap.Uint32("keyspace-group-id", group.ID))

pkg/tso/tso.go

@@ -71,6 +71,8 @@ type timestampOracle struct {
 	tsoMux *tsoObject
 	// last timestamp window stored in etcd
 	lastSavedTime atomic.Value // stored as time.Time
+	// checkTSOPrimary fences PD embedded TSO writes when a TSO microservice primary exists.
+	checkTSOPrimary func() bool
 
 	// pre-initialized metrics
 	metrics *tsoMetrics
@@ -87,7 +89,11 @@ func (t *timestampOracle) getStorageTimeout() time.Duration {
 func (t *timestampOracle) saveTimestamp(ts time.Time) error {
 	ctx, cancel := context.WithTimeout(context.Background(), t.getStorageTimeout())
 	defer cancel()
-	return t.storage.SaveTimestamp(ctx, t.keyspaceGroupID, ts, t.member.GetLeadership())
+	checkTSOPrimary := false
+	if t.checkTSOPrimary != nil {
+		checkTSOPrimary = t.checkTSOPrimary()
+	}
+	return t.storage.SaveTimestamp(ctx, t.keyspaceGroupID, ts, t.member.GetLeadership(), checkTSOPrimary)
 }
 
 // setTSOCondition is the condition for updating physical time.

server/api/member.go

@@ -81,7 +81,6 @@ func getMembers(svr *server.Server, forceRefresh bool) (*pdpb.GetMembersResponse
 	if members.GetHeader().GetError() != nil {
 		return nil, errors.WithStack(errors.New(members.GetHeader().GetError().String()))
 	}
-
 	for _, m := range members.GetMembers() {
 		var e error
 		m.BinaryVersion, e = svr.GetMember().GetMemberBinaryVersion(m.GetMemberId())