*: refactor dynamic tso switch

cmd/pd-server/main.go

@@ -52,8 +52,8 @@ import (
 )
 
 const (
-	apiMode        = "api"
-	tsoMode        = "tso"
+	// Only serverless use this variable, we can use it to determine whether the PD supports multiple timelines.
+	// The name is a little misleading, but it's kept for backward compatibility.
 	serviceModeEnv = "PD_SERVICE_MODE"
 )
 
@@ -89,7 +89,7 @@ func NewServiceCommand() *cobra.Command {
 // NewTSOServiceCommand returns the tso service command.
 func NewTSOServiceCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   tsoMode,
+		Use:   "tso",
 		Short: "Run the TSO service",
 		Run:   tso.CreateServerWrapper,
 	}
@@ -129,11 +129,12 @@ func NewSchedulingServiceCommand() *cobra.Command {
 }
 
 // NewPDServiceCommand returns the PD service command.
+// We can use pd directly. This command is kept for backward compatibility.
 func NewPDServiceCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   apiMode,
+		Use:   "api",
 		Short: "Run the PD service",
-		Run:   createPDServiceWrapper,
+		Run:   createServerWrapper,
 	}
 	addFlags(cmd)
 	return cmd
@@ -160,20 +161,7 @@ func addFlags(cmd *cobra.Command) {
 	cmd.Flags().BoolP("force-new-cluster", "", false, "force to create a new one-member cluster")
 }
 
-func createPDServiceWrapper(cmd *cobra.Command, args []string) {
-	start(cmd, args, cmd.CalledAs())
-}
-
 func createServerWrapper(cmd *cobra.Command, args []string) {
-	mode := os.Getenv(serviceModeEnv)
-	if len(mode) != 0 && strings.ToLower(mode) == apiMode {
-		start(cmd, args, apiMode)
-	} else {
-		start(cmd, args)
-	}
-}
-
-func start(cmd *cobra.Command, args []string, services ...string) {
 	schedulers.Register()
 	cfg := config.NewConfig()
 	flagSet := cmd.Flags()
@@ -183,6 +171,10 @@ func start(cmd *cobra.Command, args []string, services ...string) {
 		return
 	}
 	err = cfg.Parse(flagSet)
+	mode := os.Getenv(serviceModeEnv)
+	if len(mode) != 0 && strings.ToLower(mode) == "api" {
+		cfg.Microservice.EnableMultiTimelines = true
+	}
 	defer logutil.LogPanic()
 
 	if err != nil {
@@ -241,7 +233,7 @@ func start(cmd *cobra.Command, args []string, services ...string) {
 		serviceBuilders = append(serviceBuilders, swaggerserver.NewHandler)
 	}
 	serviceBuilders = append(serviceBuilders, dashboard.GetServiceBuilders()...)
-	svr, err := server.CreateServer(ctx, cfg, services, serviceBuilders...)
+	svr, err := server.CreateServer(ctx, cfg, serviceBuilders...)
 	if err != nil {
 		log.Fatal("create server failed", errs.ZapError(err))
 	}

pkg/keyspace/keyspace_test.go

@@ -38,9 +38,11 @@ import (
 )
 
 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 {
@@ -109,8 +111,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,
 			IsPreAlloc: true, // skip wait region split

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

@@ -85,18 +85,18 @@
 		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/utils/apiutil/serverapi/middleware.go

@@ -116,7 +116,7 @@ func MicroserviceRedirectRule(matchPath, targetPath, targetServiceName string,
 }
 
 func (h *redirector) matchMicroserviceRedirectRules(r *http.Request) (bool, string) {
-	if !h.s.IsKeyspaceGroupEnabled() {
+	if !h.s.IsServiceIndependent(constant.TSOServiceName) && !h.s.IsServiceIndependent(constant.SchedulingServiceName) {
 		return false, ""
 	}
 	if len(h.microserviceRedirectRules) == 0 {
@@ -145,7 +145,8 @@ func (h *redirector) matchMicroserviceRedirectRules(r *http.Request) (bool, stri
 			addr, ok := h.s.GetServicePrimaryAddr(r.Context(), rule.targetServiceName)
 			if !ok || addr == "" {
 				log.Warn("failed to get the service primary addr when trying to match redirect rules",
-					zap.String("path", r.URL.Path))
+					zap.String("path", r.URL.Path), zap.String("addr", addr),
+					zap.String("target", rule.targetServiceName))
 				return true, ""
 			}
 			// If the URL contains escaped characters, use RawPath instead of Path

pkg/utils/keypath/key_path.go

@@ -321,7 +321,7 @@ func svcRootPath(svcName string) string {
 	return path.Join(constant.MicroserviceRootPath, c, svcName)
 }
 
-// LegacyRootPath returns the root path of legacy pd service.
+// LegacyRootPath returns the root path of legacy PD.
 // Path: /pd/{cluster_id}
 func LegacyRootPath() string {
 	return path.Join(pdRootPath, strconv.FormatUint(ClusterID(), 10))

server/api/member.go

@@ -72,7 +72,6 @@ func getMembers(svr *server.Server) (*pdpb.GetMembersResponse, error) {
 	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())

server/api/server_test.go

@@ -98,7 +98,7 @@ func mustNewCluster(re *require.Assertions, num int, opts ...func(cfg *config.Co
 			for _, opt := range opts {
 				opt(cfg)
 			}
-			s, err := server.CreateServer(ctx, cfg, nil, NewHandler)
+			s, err := server.CreateServer(ctx, cfg, NewHandler)
 			re.NoError(err)
 			err = s.Run()
 			re.NoError(err)

server/api/version_test.go

@@ -58,7 +58,7 @@ func TestGetVersion(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	ch := make(chan *server.Server)
 	go func(cfg *config.Config) {
-		s, err := server.CreateServer(ctx, cfg, nil, NewHandler)
+		s, err := server.CreateServer(ctx, cfg, NewHandler)
 		re.NoError(err)
 		re.NoError(failpoint.Enable("github.com/tikv/pd/server/memberNil", `return(true)`))
 		reqCh <- struct{}{}

server/apiv2/handlers/micro_service.go

@@ -39,17 +39,16 @@
 // @Router   /ms/members/{service} [get]
 func GetMembers(c *gin.Context) {
 	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
-	if !svr.IsKeyspaceGroupEnabled() {
-		c.AbortWithStatusJSON(http.StatusNotFound, "not support microservice")
-		return
-	}
-
 	if service := c.Param("service"); len(service) > 0 {
 		entries, err := discovery.GetMSMembers(service, svr.GetClient())
 		if err != nil {
 			c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
 			return
 		}
+		if len(entries) == 0 {
+			c.AbortWithStatusJSON(http.StatusNotFound, "no members found")
+			return
+		}
 		c.IndentedJSON(http.StatusOK, entries)
 		return
 	}
@@ -65,13 +64,12 @@
 // @Router   /ms/primary/{service} [get]
 func GetPrimary(c *gin.Context) {
 	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
-	if !svr.IsKeyspaceGroupEnabled() {
-		c.AbortWithStatusJSON(http.StatusNotFound, "not support microservice")
-		return
-	}
-
 	if service := c.Param("service"); len(service) > 0 {
-		addr, _ := svr.GetServicePrimaryAddr(c.Request.Context(), service)
+		addr, exist := svr.GetServicePrimaryAddr(c.Request.Context(), service)
+		if !exist {
+			c.AbortWithStatusJSON(http.StatusNotFound, "no primary found")
+			return
+		}
 		c.IndentedJSON(http.StatusOK, addr)
 		return
 	}

server/cluster/cluster.go

@@ -131,7 +131,6 @@
 	GetMembers() ([]*pdpb.Member, error)
 	ReplicateFileToMember(ctx context.Context, member *pdpb.Member, name string, data []byte) error
 	GetKeyspaceGroupManager() *keyspace.GroupManager
-	IsKeyspaceGroupEnabled() bool
 	GetSafePointV2Manager() *gc.SafePointV2Manager
 }
 
@@ -156,12 +155,11 @@
 	etcdClient *clientv3.Client
 	httpClient *http.Client
 
-	running                bool
-	isKeyspaceGroupEnabled bool
-	meta                   *metapb.Cluster
-	storage                storage.Storage
-	minResolvedTS          atomic.Value // Store as uint64
-	externalTS             atomic.Value // Store as uint64
+	running       bool
+	meta          *metapb.Cluster
+	storage       storage.Storage
+	minResolvedTS atomic.Value // Store as uint64
+	externalTS    atomic.Value // Store as uint64
 
 	// Keep the previous store limit settings when removing a store.
 	prevStoreLimit map[uint64]map[storelimit.Type]float64
@@ -325,7 +323,6 @@
 		log.Warn("raft cluster has already been started")
 		return nil
 	}
-	c.isKeyspaceGroupEnabled = s.IsKeyspaceGroupEnabled()
 	err = c.InitCluster(s.GetAllocator(), s.GetPersistOptions(), s.GetHBStreams(), s.GetKeyspaceGroupManager())
 	if err != nil {
 		return err
@@ -376,14 +373,15 @@
 	c.loadExternalTS()
 	c.loadMinResolvedTS()
 
-	if c.isKeyspaceGroupEnabled {
-		// bootstrap keyspace group manager after starting other parts successfully.
-		// This order avoids a stuck goroutine in keyspaceGroupManager when it fails to create raftcluster.
+	// bootstrap keyspace group manager after starting other parts successfully.
+	// This order avoids a stuck goroutine in keyspaceGroupManager when it fails to create raftcluster.
+	if c.keyspaceGroupManager != nil {
 		err = c.keyspaceGroupManager.Bootstrap(c.ctx)
 		if err != nil {
 			return err
 		}
 	}
+
 	c.checkSchedulingService()
 	c.wg.Add(9)
 	go c.runServiceCheckJob()
@@ -404,53 +402,68 @@
 }
 
 func (c *RaftCluster) checkSchedulingService() {
-	if c.isKeyspaceGroupEnabled {
-		servers, err := discovery.Discover(c.etcdClient, constant.SchedulingServiceName)
-		if c.opt.GetMicroserviceConfig().IsSchedulingFallbackEnabled() && (err != nil || len(servers) == 0) {
-			c.startSchedulingJobs(c, c.hbstreams)
-			c.UnsetServiceIndependent(constant.SchedulingServiceName)
-		} else {
-			if c.stopSchedulingJobs() || c.coordinator == nil {
-				c.initCoordinator(c.ctx, c, c.hbstreams)
-			}
-			if !c.IsServiceIndependent(constant.SchedulingServiceName) {
-				c.SetServiceIndependent(constant.SchedulingServiceName)
-			}
-		}
-	} else {
+	servers, err := discovery.Discover(c.etcdClient, constant.SchedulingServiceName)
+	if c.opt.GetMicroserviceConfig().IsSchedulingFallbackEnabled() && (err != nil || len(servers) == 0) {
 		c.startSchedulingJobs(c, c.hbstreams)
 		c.UnsetServiceIndependent(constant.SchedulingServiceName)
+	} else {
+		if c.stopSchedulingJobs() || c.coordinator == nil {
+			c.initCoordinator(c.ctx, c, c.hbstreams)
+		}
+		if !c.IsServiceIndependent(constant.SchedulingServiceName) {
+			c.SetServiceIndependent(constant.SchedulingServiceName)
+		}
 	}
 }
 
 // checkTSOService checks the TSO service.
+// In non-serverless env:
+//  1. If the dynamic switching is disabled, it will switch to the internal TSO service.
+//  2. If the dynamic switching is enabled, it will check the external TSO service.
+//     If the external TSO service is available, it will switch to the external TSO service.
+//     If the external TSO service is unavailable, it will switch to the internal TSO service.
+//
+// In serverless env, we don't allow dynamic switching.
+// Whether we use the internal TSO service or the external TSO service is determined by the `IsMultiTimelinesEnabled`.
 func (c *RaftCluster) checkTSOService() {
-	if c.isKeyspaceGroupEnabled {
-		if c.opt.GetMicroserviceConfig().IsTSODynamicSwitchingEnabled() {
-			servers, err := discovery.Discover(c.etcdClient, constant.TSOServiceName)
-			if err != nil || len(servers) == 0 {
-				if err := c.startTSOJobsIfNeeded(); err != nil {
-					log.Error("failed to start TSO jobs", errs.ZapError(err))
-					return
-				}
-				if c.IsServiceIndependent(constant.TSOServiceName) {
-					log.Info("TSO is provided by PD")
-					c.UnsetServiceIndependent(constant.TSOServiceName)
-				}
-			} else {
-				c.stopTSOJobsIfNeeded()
-				if !c.IsServiceIndependent(constant.TSOServiceName) {
-					log.Info("TSO is provided by TSO server")
-					c.SetServiceIndependent(constant.TSOServiceName)
-				}
-			}
+	if c.opt.GetMicroserviceConfig().IsMultiTimelinesEnabled() {
+		return
+	}
+	if !c.opt.GetMicroserviceConfig().IsTSODynamicSwitchingEnabled() {
+		if err := c.switchToInternalTSO(); err != nil {
+			log.Error("failed to start TSO jobs", errs.ZapError(err))
+			return
 		}
 		return
 	}
 
+	servers, err := discovery.Discover(c.etcdClient, constant.TSOServiceName)
+	if err != nil || len(servers) == 0 {
+		if err := c.switchToInternalTSO(); err != nil {
+			log.Error("failed to switch to internal TSO", errs.ZapError(err))
+			return
+		}
+	} else if len(servers) > 0 {
+		c.switchToExternalTSO()
+	}
+}
+
+func (c *RaftCluster) switchToInternalTSO() error {
 	if err := c.startTSOJobsIfNeeded(); err != nil {
-		log.Error("failed to start TSO jobs", errs.ZapError(err))
-		return
+		return err
+	}
+	if c.IsServiceIndependent(constant.TSOServiceName) {
+		c.UnsetServiceIndependent(constant.TSOServiceName)
+		log.Info("successfully switched to internal TSO")
+	}
+	return nil
+}
+
+func (c *RaftCluster) switchToExternalTSO() {
+	c.stopTSOJobsIfNeeded()
+	if !c.IsServiceIndependent(constant.TSOServiceName) {
+		c.SetServiceIndependent(constant.TSOServiceName)
+		log.Info("successfully switched to external TSO")
 	}
 }