cmd, server: add env to manage the service mode (#7018)

ref #5766

Add env `PD_SERVICE_MODE` to manage the service mode.

Signed-off-by: Ryan Leung <rleungx@gmail.com>

Co-authored-by: Ryan Leung <rleungx@gmail.com>

cmd/pd-server/main.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
 
 	grpcprometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
@@ -48,6 +49,13 @@ import (
 	_ "github.com/tikv/pd/pkg/mcs/tso/server/install"
 )
 
+const (
+	apiMode        = "api"
+	tsoMode        = "tso"
+	rmMode         = "resource-manager"
+	serviceModeEnv = "PD_SERVICE_MODE"
+)
+
 func main() {
 	rootCmd := &cobra.Command{
 		Use:   "pd-server",
@@ -81,7 +89,7 @@ func NewServiceCommand() *cobra.Command {
 // NewTSOServiceCommand returns the tso service command.
 func NewTSOServiceCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "tso",
+		Use:   tsoMode,
 		Short: "Run the TSO service",
 		Run:   tso.CreateServerWrapper,
 	}
@@ -121,7 +129,7 @@ func NewSchedulingServiceCommand() *cobra.Command {
 // NewResourceManagerServiceCommand returns the resource manager service command.
 func NewResourceManagerServiceCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "resource-manager",
+		Use:   rmMode,
 		Short: "Run the resource manager service",
 		Run:   resource_manager.CreateServerWrapper,
 	}
@@ -141,7 +149,7 @@ func NewResourceManagerServiceCommand() *cobra.Command {
 // NewAPIServiceCommand returns the API service command.
 func NewAPIServiceCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "api",
+		Use:   apiMode,
 		Short: "Run the API service",
 		Run:   createAPIServerWrapper,
 	}
@@ -175,7 +183,12 @@ func createAPIServerWrapper(cmd *cobra.Command, args []string) {
 }
 
 func createServerWrapper(cmd *cobra.Command, args []string) {
-	start(cmd, args)
+	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) {

server/server.go

@@ -237,10 +237,10 @@ type HandlerBuilder func(context.Context, *Server) (http.Handler, apiutil.APISer
 // CreateServer creates the UNINITIALIZED pd server with given configuration.
 func CreateServer(ctx context.Context, cfg *config.Config, services []string, legacyServiceBuilders ...HandlerBuilder) (*Server, error) {
 	var mode string
-	if len(services) == 0 {
-		mode = PDMode
-	} else {
+	if len(services) != 0 {
 		mode = APIServiceMode
+	} else {
+		mode = PDMode
 	}
 	log.Info(fmt.Sprintf("%s config", mode), zap.Reflect("config", cfg))
 	serviceMiddlewareCfg := config.NewServiceMiddlewareConfig()