Merge branch 'master' into split-scheduler-test

pkg/basicserver/basic_server.go

@@ -44,5 +44,5 @@ type Server interface {
 	// IsServing returns whether the server is the leader, if there is embedded etcd, or the primary otherwise.
 	IsServing() bool
 	// AddServiceReadyCallback adds callbacks when the server becomes the leader, if there is embedded etcd, or the primary otherwise.
-	AddServiceReadyCallback(callbacks ...func(context.Context))
+	AddServiceReadyCallback(callbacks ...func(context.Context) error)
 }

pkg/mcs/resourcemanager/server/manager.go

@@ -101,7 +101,7 @@ func (m *Manager) GetBasicServer() bs.Server {
 }
 
 // Init initializes the resource group manager.
-func (m *Manager) Init(ctx context.Context) {
+func (m *Manager) Init(ctx context.Context) error {
 	// Todo: If we can modify following configs in the future, we should reload these configs.
 	// Store the controller config into the storage.
 	m.storage.SaveControllerConfig(m.controllerConfig)
@@ -156,6 +156,7 @@ func (m *Manager) Init(ctx context.Context) {
 		m.persistLoop(ctx)
 	}()
 	log.Info("resource group manager finishes initialization")
+	return nil
 }
 
 // AddResourceGroup puts a resource group.

pkg/mcs/resourcemanager/server/server.go

@@ -71,7 +71,7 @@ type Server struct {
 	service *Service
 
 	// primaryCallbacks will be called after the server becomes leader.
-	primaryCallbacks []func(context.Context)
+	primaryCallbacks []func(context.Context) error
 
 	serviceRegister *discovery.ServiceRegister
 }
@@ -232,7 +232,7 @@ func (s *Server) IsClosed() bool {
 }
 
 // AddServiceReadyCallback adds callbacks when the server becomes the leader, if there is embedded etcd, or the primary otherwise.
-func (s *Server) AddServiceReadyCallback(callbacks ...func(context.Context)) {
+func (s *Server) AddServiceReadyCallback(callbacks ...func(context.Context) error) {
 	s.primaryCallbacks = append(s.primaryCallbacks, callbacks...)
 }
 

pkg/mcs/scheduling/server/cluster.go

@@ -143,19 +143,27 @@ func (c *Cluster) GetStoreConfig() sc.StoreConfigProvider { return c.persistConf
 
 // AllocID allocates a new ID.
 func (c *Cluster) AllocID() (uint64, error) {
-	cli := c.apiServerLeader.Load().(pdpb.PDClient)
-	if cli == nil {
-		c.checkMembershipCh <- struct{}{}
-		return 0, errors.New("API server leader is not found")
+	client, err := c.getAPIServerLeaderClient()
+	if err != nil {
+		return 0, err
 	}
-	resp, err := cli.AllocID(c.ctx, &pdpb.AllocIDRequest{Header: &pdpb.RequestHeader{ClusterId: c.clusterID}})
+	resp, err := client.AllocID(c.ctx, &pdpb.AllocIDRequest{Header: &pdpb.RequestHeader{ClusterId: c.clusterID}})
 	if err != nil {
 		c.checkMembershipCh <- struct{}{}
 		return 0, err
 	}
 	return resp.GetId(), nil
 }
 
+func (c *Cluster) getAPIServerLeaderClient() (pdpb.PDClient, error) {
+	cli := c.apiServerLeader.Load()
+	if cli == nil {
+		c.checkMembershipCh <- struct{}{}
+		return nil, errors.New("API server leader is not found")
+	}
+	return cli.(pdpb.PDClient), nil
+}
+
 // SwitchAPIServerLeader switches the API server leader.
 func (c *Cluster) SwitchAPIServerLeader(new pdpb.PDClient) bool {
 	old := c.apiServerLeader.Load()

pkg/mcs/scheduling/server/config/config.go

@@ -34,6 +34,7 @@ import (
 	"github.com/tikv/pd/pkg/core/storelimit"
 	"github.com/tikv/pd/pkg/mcs/utils"
 	sc "github.com/tikv/pd/pkg/schedule/config"
+	"github.com/tikv/pd/pkg/slice"
 	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/utils/configutil"
 	"github.com/tikv/pd/pkg/utils/grpcutil"
@@ -239,6 +240,18 @@ func (o *PersistConfig) SetScheduleConfig(cfg *sc.ScheduleConfig) {
 	o.schedule.Store(cfg)
 }
 
+// AdjustScheduleCfg adjusts the schedule config.
+func (o *PersistConfig) AdjustScheduleCfg(scheduleCfg *sc.ScheduleConfig) {
+	// In case we add new default schedulers.
+	for _, ps := range sc.DefaultSchedulers {
+		if slice.NoneOf(scheduleCfg.Schedulers, func(i int) bool {
+			return scheduleCfg.Schedulers[i].Type == ps.Type
+		}) {
+			scheduleCfg.Schedulers = append(scheduleCfg.Schedulers, ps)
+		}
+	}
+}
+
 // GetReplicationConfig returns replication configurations.
 func (o *PersistConfig) GetReplicationConfig() *sc.ReplicationConfig {
 	return o.replication.Load().(*sc.ReplicationConfig)

pkg/mcs/scheduling/server/config/watcher.go

@@ -94,6 +94,7 @@ func (cw *Watcher) initializeConfigWatcher() error {
 				zap.String("event-kv-key", string(kv.Key)), zap.Error(err))
 			return err
 		}
+		cw.AdjustScheduleCfg(&cfg.Schedule)
 		cw.SetClusterVersion(&cfg.ClusterVersion)
 		cw.SetScheduleConfig(&cfg.Schedule)
 		cw.SetReplicationConfig(&cfg.Replication)

pkg/mcs/scheduling/server/server.go

@@ -87,7 +87,8 @@ type Server struct {
 	checkMembershipCh chan struct{}
 
 	// primaryCallbacks will be called after the server becomes leader.
-	primaryCallbacks []func(context.Context)
+	primaryCallbacks     []func(context.Context) error
+	primaryExitCallbacks []func()
 
 	// for service registry
 	serviceID       *discovery.ServiceRegistryEntry
@@ -164,6 +165,9 @@ func (s *Server) updateAPIServerMemberLoop() {
 		case <-ticker.C:
 		case <-s.checkMembershipCh:
 		}
+		if !s.IsServing() {
+			continue
+		}
 		members, err := s.GetClient().MemberList(ctx)
 		if err != nil {
 			log.Warn("failed to list members", errs.ZapError(err))
@@ -247,9 +251,16 @@ func (s *Server) campaignLeader() {
 
 	log.Info("triggering the primary callback functions")
 	for _, cb := range s.primaryCallbacks {
-		cb(ctx)
+		if err := cb(ctx); err != nil {
+			log.Error("failed to trigger the primary callback functions", errs.ZapError(err))
+			return
+		}
 	}
-
+	defer func() {
+		for _, cb := range s.primaryExitCallbacks {
+			cb()
+		}
+	}()
 	s.participant.EnableLeader()
 	log.Info("scheduling primary is ready to serve", zap.String("scheduling-primary-name", s.participant.Name()))
 
@@ -283,10 +294,6 @@ func (s *Server) Close() {
 	utils.StopHTTPServer(s)
 	utils.StopGRPCServer(s)
 	s.GetListener().Close()
-	s.GetCoordinator().Stop()
-	s.ruleWatcher.Close()
-	s.configWatcher.Close()
-	s.metaWatcher.Close()
 	s.serverLoopCancel()
 	s.serverLoopWg.Wait()
 
@@ -313,10 +320,15 @@ func (s *Server) IsClosed() bool {
 }
 
 // AddServiceReadyCallback adds callbacks when the server becomes the leader, if there is embedded etcd, or the primary otherwise.
-func (s *Server) AddServiceReadyCallback(callbacks ...func(context.Context)) {
+func (s *Server) AddServiceReadyCallback(callbacks ...func(context.Context) error) {
 	s.primaryCallbacks = append(s.primaryCallbacks, callbacks...)
 }
 
+// AddServiceExitCallback adds callbacks when the server becomes the leader, if there is embedded etcd, or the primary otherwise.
+func (s *Server) AddServiceExitCallback(callbacks ...func()) {
+	s.primaryExitCallbacks = append(s.primaryExitCallbacks, callbacks...)
+}
+
 // GetTLSConfig gets the security config.
 func (s *Server) GetTLSConfig() *grpcutil.TLSConfig {
 	return &s.cfg.Security.TLSConfig
@@ -381,20 +393,10 @@ func (s *Server) startServer() (err error) {
 		ListenUrls: []string{s.cfg.AdvertiseListenAddr},
 	}
 	s.participant.InitInfo(p, endpoint.SchedulingSvcRootPath(s.clusterID), utils.PrimaryKey, "primary election")
-	s.basicCluster = core.NewBasicCluster()
-	err = s.startWatcher()
-	if err != nil {
-		return err
-	}
-	s.storage = endpoint.NewStorageEndpoint(
-		kv.NewEtcdKVBase(s.GetClient(), endpoint.PDRootPath(s.clusterID)), nil)
-	s.hbStreams = hbstream.NewHeartbeatStreams(s.Context(), s.clusterID, s.basicCluster)
-	s.cluster, err = NewCluster(s.Context(), s.persistConfig, s.storage, s.basicCluster, s.hbStreams, s.clusterID, s.checkMembershipCh)
-	if err != nil {
-		return err
-	}
 
 	s.service = &Service{Server: s}
+	s.AddServiceReadyCallback(s.startCluster)
+	s.AddServiceExitCallback(s.stopCluster)
 	if err := s.InitListener(s.GetTLSConfig(), s.cfg.ListenAddr); err != nil {
 		return err
 	}
@@ -406,7 +408,6 @@ func (s *Server) startServer() (err error) {
 	go utils.StartGRPCAndHTTPServers(s, serverReadyChan, s.GetListener())
 	s.checkMembershipCh <- struct{}{}
 	<-serverReadyChan
-	go s.GetCoordinator().RunUntilStop()
 
 	// Run callbacks
 	log.Info("triggering the start callback functions")
@@ -429,6 +430,29 @@ func (s *Server) startServer() (err error) {
 	return nil
 }
 
+func (s *Server) startCluster(context.Context) error {
+	s.basicCluster = core.NewBasicCluster()
+	err := s.startWatcher()
+	if err != nil {
+		return err
+	}
+	s.storage = endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
+	s.hbStreams = hbstream.NewHeartbeatStreams(s.Context(), s.clusterID, s.basicCluster)
+	s.cluster, err = NewCluster(s.Context(), s.persistConfig, s.storage, s.basicCluster, s.hbStreams, s.clusterID, s.checkMembershipCh)
+	if err != nil {
+		return err
+	}
+	go s.GetCoordinator().RunUntilStop()
+	return nil
+}
+
+func (s *Server) stopCluster() {
+	s.GetCoordinator().Stop()
+	s.ruleWatcher.Close()
+	s.configWatcher.Close()
+	s.metaWatcher.Close()
+}
+
 func (s *Server) startWatcher() (err error) {
 	s.metaWatcher, err = meta.NewWatcher(s.Context(), s.GetClient(), s.clusterID, s.basicCluster)
 	if err != nil {

pkg/mcs/tso/server/server.go

@@ -236,7 +236,7 @@ func (s *Server) ResignPrimary(keyspaceID, keyspaceGroupID uint32) error {
 
 // AddServiceReadyCallback implements basicserver.
 // It adds callbacks when it's ready for providing tso service.
-func (s *Server) AddServiceReadyCallback(callbacks ...func(context.Context)) {
+func (s *Server) AddServiceReadyCallback(callbacks ...func(context.Context) error) {
 	// Do nothing here. The primary of each keyspace group assigned to this host
 	// will respond to the requests accordingly.
 }

server/api/min_resolved_ts.go

@@ -53,7 +53,7 @@ type minResolvedTS struct {
 // @Failure      500    {string}  string  "PD server failed to proceed the request."
 // @Router   /min-resolved-ts/{store_id} [get]
 func (h *minResolvedTSHandler) GetStoreMinResolvedTS(w http.ResponseWriter, r *http.Request) {
-	c := h.svr.GetRaftCluster()
+	c := getCluster(r)
 	idStr := mux.Vars(r)["store_id"]
 	storeID, err := strconv.ParseUint(idStr, 10, 64)
 	if err != nil {
@@ -84,7 +84,7 @@ func (h *minResolvedTSHandler) GetStoreMinResolvedTS(w http.ResponseWriter, r *h
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router       /min-resolved-ts [get]
 func (h *minResolvedTSHandler) GetMinResolvedTS(w http.ResponseWriter, r *http.Request) {
-	c := h.svr.GetRaftCluster()
+	c := getCluster(r)
 	scopeMinResolvedTS := c.GetMinResolvedTS()
 	persistInterval := c.GetPDServerConfig().MinResolvedTSPersistenceInterval
 

server/grpc_service.go

@@ -1600,7 +1600,6 @@ func (s *GrpcServer) ReportBatchSplit(ctx context.Context, request *pdpb.ReportB
 	if rc == nil {
 		return &pdpb.ReportBatchSplitResponse{Header: s.notBootstrappedHeader()}, nil
 	}
-
 	_, err := rc.HandleBatchReportSplit(request)
 	if err != nil {
 		return &pdpb.ReportBatchSplitResponse{
@@ -2089,6 +2088,9 @@ func (s *GrpcServer) SplitAndScatterRegions(ctx context.Context, request *pdpb.S
 		return rsp.(*pdpb.SplitAndScatterRegionsResponse), err
 	}
 	rc := s.GetRaftCluster()
+	if rc == nil {
+		return &pdpb.SplitAndScatterRegionsResponse{Header: s.notBootstrappedHeader()}, nil
+	}
 	splitFinishedPercentage, newRegionIDs := rc.GetRegionSplitter().SplitRegions(ctx, request.GetSplitKeys(), int(request.GetRetryLimit()))
 	scatterFinishedPercentage, err := scatterRegions(rc, newRegionIDs, request.GetGroup(), int(request.GetRetryLimit()), false)
 	if err != nil {

server/handler.go

@@ -961,14 +961,20 @@ func (h *Handler) ResetTS(ts uint64, ignoreSmaller, skipUpperBoundCheck bool, _
 
 // SetStoreLimitScene sets the limit values for different scenes
 func (h *Handler) SetStoreLimitScene(scene *storelimit.Scene, limitType storelimit.Type) {
-	cluster := h.s.GetRaftCluster()
-	cluster.GetStoreLimiter().ReplaceStoreLimitScene(scene, limitType)
+	rc := h.s.GetRaftCluster()
+	if rc == nil {
+		return
+	}
+	rc.GetStoreLimiter().ReplaceStoreLimitScene(scene, limitType)
 }
 
 // GetStoreLimitScene returns the limit values for different scenes
 func (h *Handler) GetStoreLimitScene(limitType storelimit.Type) *storelimit.Scene {
-	cluster := h.s.GetRaftCluster()
-	return cluster.GetStoreLimiter().StoreLimitScene(limitType)
+	rc := h.s.GetRaftCluster()
+	if rc == nil {
+		return nil
+	}
+	return rc.GetStoreLimiter().StoreLimitScene(limitType)
 }
 
 // GetProgressByID returns the progress details for a given store ID.