mcs/scheduling: register the scheduler handler in API service mode (#7082)

ref #5839

To ensure that the HTTP API handler could be initialized properly for each scheduler in the API service mode,
this PR updates the scheduler controller and PD server handler to support initialize the HTTP handler only.

Signed-off-by: JmPotato <ghzpotato@gmail.com>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>

Author: JmPotato

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

@@ -19,7 +19,6 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
-	"sync"
 	"sync/atomic"
 	"time"
 	"unsafe"
@@ -204,8 +203,6 @@ type PersistConfig struct {
 	schedule       atomic.Value
 	replication    atomic.Value
 	storeConfig    atomic.Value
-	// Store the respective configurations for different schedulers.
-	schedulerConfig sync.Map
 }
 
 // NewPersistConfig creates a new PersistConfig instance.
@@ -275,24 +272,6 @@ func (o *PersistConfig) GetStoreConfig() *sc.StoreConfig {
 	return o.storeConfig.Load().(*sc.StoreConfig)
 }
 
-// SetSchedulerConfig sets the scheduler configuration with the given name.
-func (o *PersistConfig) SetSchedulerConfig(name, data string) {
-	o.schedulerConfig.Store(name, data)
-}
-
-// RemoveSchedulerConfig removes the scheduler configuration with the given name.
-func (o *PersistConfig) RemoveSchedulerConfig(name string) {
-	o.schedulerConfig.Delete(name)
-}
-
-// GetSchedulerConfig returns the scheduler configuration with the given name.
-func (o *PersistConfig) GetSchedulerConfig(name string) string {
-	if v, ok := o.schedulerConfig.Load(name); ok {
-		return v.(string)
-	}
-	return ""
-}
-
 // GetMaxReplicas returns the max replicas.
 func (o *PersistConfig) GetMaxReplicas() int {
 	return int(o.GetReplicationConfig().MaxReplicas)

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

@@ -23,6 +23,7 @@ import (
 	"github.com/coreos/go-semver/semver"
 	"github.com/pingcap/log"
 	sc "github.com/tikv/pd/pkg/schedule/config"
+	"github.com/tikv/pd/pkg/storage"
 	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/utils/etcdutil"
 	"go.etcd.io/etcd/clientv3"
@@ -50,6 +51,9 @@ type Watcher struct {
 	schedulerConfigWatcher *etcdutil.LoopWatcher
 
 	*PersistConfig
+	// Some data, like the scheduler configs, should be loaded into the storage
+	// to make sure the coordinator could access them correctly.
+	storage storage.Storage
 }
 
 type persistedConfig struct {
@@ -65,6 +69,7 @@ func NewWatcher(
 	etcdClient *clientv3.Client,
 	clusterID uint64,
 	persistConfig *PersistConfig,
+	storage storage.Storage,
 ) (*Watcher, error) {
 	ctx, cancel := context.WithCancel(ctx)
 	cw := &Watcher{
@@ -74,6 +79,7 @@ func NewWatcher(
 		schedulerConfigPathPrefix: endpoint.SchedulerConfigPathPrefix(clusterID),
 		etcdClient:                etcdClient,
 		PersistConfig:             persistConfig,
+		storage:                   storage,
 	}
 	err := cw.initializeConfigWatcher()
 	if err != nil {
@@ -120,15 +126,15 @@ func (cw *Watcher) initializeConfigWatcher() error {
 func (cw *Watcher) initializeSchedulerConfigWatcher() error {
 	prefixToTrim := cw.schedulerConfigPathPrefix + "/"
 	putFn := func(kv *mvccpb.KeyValue) error {
-		cw.SetSchedulerConfig(
+		return cw.storage.SaveScheduleConfig(
 			strings.TrimPrefix(string(kv.Key), prefixToTrim),
-			string(kv.Value),
+			kv.Value,
 		)
-		return nil
 	}
 	deleteFn := func(kv *mvccpb.KeyValue) error {
-		cw.RemoveSchedulerConfig(strings.TrimPrefix(string(kv.Key), prefixToTrim))
-		return nil
+		return cw.storage.RemoveScheduleConfig(
+			strings.TrimPrefix(string(kv.Key), prefixToTrim),
+		)
 	}
 	postEventFn := func() error {
 		return nil

pkg/mcs/scheduling/server/server.go

@@ -432,11 +432,11 @@ func (s *Server) startServer() (err error) {
 
 func (s *Server) startCluster(context.Context) error {
 	s.basicCluster = core.NewBasicCluster()
+	s.storage = endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
 	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 {
@@ -458,7 +458,7 @@ func (s *Server) startWatcher() (err error) {
 	if err != nil {
 		return err
 	}
-	s.configWatcher, err = config.NewWatcher(s.Context(), s.GetClient(), s.clusterID, s.persistConfig)
+	s.configWatcher, err = config.NewWatcher(s.Context(), s.GetClient(), s.clusterID, s.persistConfig, s.storage)
 	if err != nil {
 		return err
 	}

pkg/schedule/coordinator.go

@@ -444,6 +444,8 @@ func (c *Coordinator) InitSchedulers(needRun bool) {
 			if err = c.schedulers.AddScheduler(s); err != nil {
 				log.Error("can not add scheduler with independent configuration", zap.String("scheduler-name", s.GetName()), zap.Strings("scheduler-args", cfg.Args), errs.ZapError(err))
 			}
+		} else if err = c.schedulers.AddSchedulerHandler(s); err != nil {
+			log.Error("can not add scheduler handler with independent configuration", zap.String("scheduler-name", s.GetName()), zap.Strings("scheduler-args", cfg.Args), errs.ZapError(err))
 		}
 	}
 
@@ -472,6 +474,8 @@ func (c *Coordinator) InitSchedulers(needRun bool) {
 				scheduleCfg.Schedulers[k] = schedulerCfg
 				k++
 			}
+		} else if err = c.schedulers.AddSchedulerHandler(s, schedulerCfg.Args...); err != nil && !errors.ErrorEqual(err, errs.ErrSchedulerExisted.FastGenByArgs()) {
+			log.Error("can not add scheduler handler", zap.String("scheduler-name", s.GetName()), zap.Strings("scheduler-args", schedulerCfg.Args), errs.ZapError(err))
 		}
 	}
 
@@ -507,6 +511,7 @@ func (c *Coordinator) LoadPlugin(pluginPath string, ch chan string) {
 		return
 	}
 	log.Info("create scheduler", zap.String("scheduler-name", s.GetName()))
+	// TODO: handle the plugin in API service mode.
 	if err = c.schedulers.AddScheduler(s); err != nil {
 		log.Error("can't add scheduler", zap.String("scheduler-name", s.GetName()), errs.ZapError(err))
 		return

pkg/schedule/schedulers/scheduler.go

@@ -91,8 +91,10 @@ func ConfigSliceDecoder(name string, args []string) ConfigDecoder {
 // CreateSchedulerFunc is for creating scheduler.
 type CreateSchedulerFunc func(opController *operator.Controller, storage endpoint.ConfigStorage, dec ConfigDecoder, removeSchedulerCb ...func(string) error) (Scheduler, error)
 
-var schedulerMap = make(map[string]CreateSchedulerFunc)
-var schedulerArgsToDecoder = make(map[string]ConfigSliceDecoderBuilder)
+var (
+	schedulerMap           = make(map[string]CreateSchedulerFunc)
+	schedulerArgsToDecoder = make(map[string]ConfigSliceDecoderBuilder)
+)
 
 // RegisterScheduler binds a scheduler creator. It should be called in init()
 // func of a package.

pkg/schedule/schedulers/scheduler_controller.go

@@ -40,22 +40,28 @@ var denySchedulersByLabelerCounter = labeler.LabelerEventCounter.WithLabelValues
 // Controller is used to manage all schedulers.
 type Controller struct {
 	sync.RWMutex
-	wg           sync.WaitGroup
-	ctx          context.Context
-	cluster      sche.SchedulerCluster
-	storage      endpoint.ConfigStorage
-	schedulers   map[string]*ScheduleController
-	opController *operator.Controller
+	wg      sync.WaitGroup
+	ctx     context.Context
+	cluster sche.SchedulerCluster
+	storage endpoint.ConfigStorage
+	// schedulers is used to manage all schedulers, which will only be initialized
+	// and used in the PD leader service mode now.
+	schedulers map[string]*ScheduleController
+	// schedulerHandlers is used to manage the HTTP handlers of schedulers,
+	// which will only be initialized and used in the API service mode now.
+	schedulerHandlers map[string]http.Handler
+	opController      *operator.Controller
 }
 
 // NewController creates a scheduler controller.
 func NewController(ctx context.Context, cluster sche.SchedulerCluster, storage endpoint.ConfigStorage, opController *operator.Controller) *Controller {
 	return &Controller{
-		ctx:          ctx,
-		cluster:      cluster,
-		storage:      storage,
-		schedulers:   make(map[string]*ScheduleController),
-		opController: opController,
+		ctx:               ctx,
+		cluster:           cluster,
+		storage:           storage,
+		schedulers:        make(map[string]*ScheduleController),
+		schedulerHandlers: make(map[string]http.Handler),
+		opController:      opController,
 	}
 }
 
@@ -86,6 +92,9 @@ func (c *Controller) GetSchedulerNames() []string {
 func (c *Controller) GetSchedulerHandlers() map[string]http.Handler {
 	c.RLock()
 	defer c.RUnlock()
+	if len(c.schedulerHandlers) > 0 {
+		return c.schedulerHandlers
+	}
 	handlers := make(map[string]http.Handler, len(c.schedulers))
 	for name, scheduler := range c.schedulers {
 		handlers[name] = scheduler.Scheduler
@@ -117,6 +126,50 @@ func (c *Controller) ResetSchedulerMetrics() {
 	schedulerStatusGauge.Reset()
 }
 
+// AddSchedulerHandler adds the HTTP handler for a scheduler.
+func (c *Controller) AddSchedulerHandler(scheduler Scheduler, args ...string) error {
+	c.Lock()
+	defer c.Unlock()
+
+	name := scheduler.GetName()
+	if _, ok := c.schedulerHandlers[name]; ok {
+		return errs.ErrSchedulerExisted.FastGenByArgs()
+	}
+
+	c.schedulerHandlers[name] = scheduler
+	c.cluster.GetSchedulerConfig().AddSchedulerCfg(scheduler.GetType(), args)
+	return nil
+}
+
+// RemoveSchedulerHandler removes the HTTP handler for a scheduler.
+func (c *Controller) RemoveSchedulerHandler(name string) error {
+	c.Lock()
+	defer c.Unlock()
+	if c.cluster == nil {
+		return errs.ErrNotBootstrapped.FastGenByArgs()
+	}
+	s, ok := c.schedulerHandlers[name]
+	if !ok {
+		return errs.ErrSchedulerNotFound.FastGenByArgs()
+	}
+
+	conf := c.cluster.GetSchedulerConfig()
+	conf.RemoveSchedulerCfg(s.(Scheduler).GetType())
+	if err := conf.Persist(c.storage); err != nil {
+		log.Error("the option can not persist scheduler config", errs.ZapError(err))
+		return err
+	}
+
+	if err := c.storage.RemoveScheduleConfig(name); err != nil {
+		log.Error("can not remove the scheduler config", errs.ZapError(err))
+		return err
+	}
+
+	delete(c.schedulerHandlers, name)
+
+	return nil
+}
+
 // AddScheduler adds a scheduler.
 func (c *Controller) AddScheduler(scheduler Scheduler, args ...string) error {
 	c.Lock()
@@ -249,8 +302,9 @@ func (c *Controller) IsSchedulerExisted(name string) (bool, error) {
 	if c.cluster == nil {
 		return false, errs.ErrNotBootstrapped.FastGenByArgs()
 	}
-	_, ok := c.schedulers[name]
-	if !ok {
+	_, existScheduler := c.schedulers[name]
+	_, existHandler := c.schedulerHandlers[name]
+	if !existScheduler && !existHandler {
 		return false, errs.ErrSchedulerNotFound.FastGenByArgs()
 	}
 	return true, nil

server/cluster/cluster.go

@@ -795,6 +795,16 @@ func (c *RaftCluster) GetSchedulerHandlers() map[string]http.Handler {
 	return c.coordinator.GetSchedulersController().GetSchedulerHandlers()
 }
 
+// AddSchedulerHandler adds a scheduler handler.
+func (c *RaftCluster) AddSchedulerHandler(scheduler schedulers.Scheduler, args ...string) error {
+	return c.coordinator.GetSchedulersController().AddSchedulerHandler(scheduler, args...)
+}
+
+// RemoveSchedulerHandler removes a scheduler handler.
+func (c *RaftCluster) RemoveSchedulerHandler(name string) error {
+	return c.coordinator.GetSchedulersController().RemoveSchedulerHandler(name)
+}
+
 // AddScheduler adds a scheduler.
 func (c *RaftCluster) AddScheduler(scheduler schedulers.Scheduler, args ...string) error {
 	return c.coordinator.GetSchedulersController().AddScheduler(scheduler, args...)

server/handler.go

@@ -236,19 +236,24 @@ func (h *Handler) AddScheduler(name string, args ...string) error {
 		return err
 	}
 	log.Info("create scheduler", zap.String("scheduler-name", s.GetName()), zap.Strings("scheduler-args", args))
-	if !h.s.IsAPIServiceMode() {
+	if h.s.IsAPIServiceMode() {
+		if err = c.AddSchedulerHandler(s, args...); err != nil {
+			log.Error("can not add scheduler handler", zap.String("scheduler-name", s.GetName()), zap.Strings("scheduler-args", args), errs.ZapError(err))
+			return err
+		}
+		log.Info("add scheduler handler successfully", zap.String("scheduler-name", name), zap.Strings("scheduler-args", args))
+	} else {
 		if err = c.AddScheduler(s, args...); err != nil {
 			log.Error("can not add scheduler", zap.String("scheduler-name", s.GetName()), zap.Strings("scheduler-args", args), errs.ZapError(err))
 			return err
 		}
-	} else {
-		c.GetSchedulerConfig().AddSchedulerCfg(s.GetType(), args)
+		log.Info("add scheduler successfully", zap.String("scheduler-name", name), zap.Strings("scheduler-args", args))
 	}
 	if err = h.opt.Persist(c.GetStorage()); err != nil {
 		log.Error("can not persist scheduler config", errs.ZapError(err))
 		return err
 	}
-	log.Info("add scheduler successfully", zap.String("scheduler-name", name), zap.Strings("scheduler-args", args))
+	log.Info("persist scheduler config successfully", zap.String("scheduler-name", name), zap.Strings("scheduler-args", args))
 	return nil
 }
 
@@ -258,24 +263,18 @@ func (h *Handler) RemoveScheduler(name string) error {
 	if err != nil {
 		return err
 	}
-	if !h.s.IsAPIServiceMode() {
+	if h.s.IsAPIServiceMode() {
+		if err = c.RemoveSchedulerHandler(name); err != nil {
+			log.Error("can not remove scheduler handler", zap.String("scheduler-name", name), errs.ZapError(err))
+		} else {
+			log.Info("remove scheduler handler successfully", zap.String("scheduler-name", name))
+		}
+	} else {
 		if err = c.RemoveScheduler(name); err != nil {
 			log.Error("can not remove scheduler", zap.String("scheduler-name", name), errs.ZapError(err))
 		} else {
 			log.Info("remove scheduler successfully", zap.String("scheduler-name", name))
 		}
-	} else {
-		conf := c.GetSchedulerConfig()
-		c.GetSchedulerConfig().RemoveSchedulerCfg(schedulers.FindSchedulerTypeByName(name))
-		if err := conf.Persist(c.GetStorage()); err != nil {
-			log.Error("the option can not persist scheduler config", errs.ZapError(err))
-			return err
-		}
-
-		if err := c.GetStorage().RemoveScheduleConfig(name); err != nil {
-			log.Error("can not remove the scheduler config", errs.ZapError(err))
-			return err
-		}
 	}
 	return err
 }