server, schedule: make config updates atomic

pkg/dashboard/adapter/manager.go

@@ -185,9 +185,7 @@ func (m *Manager) setNewAddress() {
 	}
 
 	// set new dashboard address
-	cfg := m.srv.GetPersistOptions().GetPDServerConfig().Clone()
-	cfg.DashboardAddress = m.members[minMemberIdx].GetClientUrls()[0]
-	if err := m.srv.SetPDServerConfig(*cfg); err != nil {
+	if err := m.srv.SetDashboardAddress(m.members[minMemberIdx].GetClientUrls()[0]); err != nil {
 		log.Warn("failed to set persist options")
 	}
 }

pkg/mcs/scheduling/server/cluster.go

@@ -107,6 +107,7 @@ type Cluster struct {
 const (
 	regionLabelGCInterval = time.Hour
 	requestTimeout        = 3 * time.Second
+	pdLeaderRetryInterval = 50 * time.Millisecond
 	requestInterval       = 10 * time.Second
 
 	// heartbeat relative const
@@ -382,12 +383,12 @@ func (c *Cluster) GetStoreConfig() sc.StoreConfigProvider { return c.persistConf
 
 // AllocID allocates new IDs.
 func (c *Cluster) AllocID(count uint32) (uint64, uint32, error) {
-	client, err := c.getPDLeaderClient()
+	ctx, cancel := context.WithTimeout(c.ctx, requestTimeout)
+	defer cancel()
+	client, err := c.getPDLeaderClient(ctx)
 	if err != nil {
 		return 0, 0, err
 	}
-	ctx, cancel := context.WithTimeout(c.ctx, requestTimeout)
-	defer cancel()
 	req := &pdpb.AllocIDRequest{Header: &pdpb.RequestHeader{ClusterId: keypath.ClusterID()}, Count: count}
 
 	failpoint.Inject("allocIDNonBatch", func() {
@@ -401,13 +402,24 @@ func (c *Cluster) AllocID(count uint32) (uint64, uint32, error) {
 	return resp.GetId(), resp.GetCount(), nil
 }
 
-func (c *Cluster) getPDLeaderClient() (pdpb.PDClient, error) {
-	cli := c.pdLeader.Load()
-	if cli == nil {
+func (c *Cluster) getPDLeaderClient(ctx context.Context) (pdpb.PDClient, error) {
+	if cli := c.pdLeader.Load(); cli != nil {
+		return cli.(pdpb.PDClient), nil
+	}
+	c.triggerMembershipCheck()
+	ticker := time.NewTicker(pdLeaderRetryInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, errors.New("pd leader is not found")
+		case <-ticker.C:
+		}
+		if cli := c.pdLeader.Load(); cli != nil {
+			return cli.(pdpb.PDClient), nil
+		}
 		c.triggerMembershipCheck()
-		return nil, errors.New("PD leader is not found")
 	}
-	return cli.(pdpb.PDClient), nil
 }
 
 func (c *Cluster) triggerMembershipCheck() {

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

@@ -19,6 +19,7 @@ import (
 	"os"
 	"reflect"
 	"strconv"
+	"sync"
 	"sync/atomic"
 	"time"
 	"unsafe"
@@ -190,7 +191,8 @@ func (c *Config) Clone() *Config {
 // PersistConfig wraps all configurations that need to persist to storage and
 // allows to access them safely.
 type PersistConfig struct {
-	ttl *cache.TTLString
+	ttl      *cache.TTLString
+	configMu sync.Mutex
 	// Store the global configuration that is related to the scheduling.
 	clusterVersion unsafe.Pointer
 	schedule       atomic.Value
@@ -232,7 +234,10 @@ func (o *PersistConfig) tryNotifySchedulersUpdating() {
 	if notifier == nil {
 		return
 	}
-	notifier <- struct{}{}
+	select {
+	case notifier <- struct{}{}:
+	default:
+	}
 }
 
 // GetClusterVersion returns the cluster version.
@@ -252,6 +257,12 @@ func (o *PersistConfig) GetScheduleConfig() *sc.ScheduleConfig {
 
 // SetScheduleConfig sets the scheduling configuration dynamically.
 func (o *PersistConfig) SetScheduleConfig(cfg *sc.ScheduleConfig) {
+	o.configMu.Lock()
+	defer o.configMu.Unlock()
+	o.storeScheduleConfig(cfg)
+}
+
+func (o *PersistConfig) storeScheduleConfig(cfg *sc.ScheduleConfig) {
 	old := o.GetScheduleConfig()
 	o.schedule.Store(cfg)
 	// The coordinator is not aware of the underlying scheduler config changes,
@@ -261,6 +272,24 @@ func (o *PersistConfig) SetScheduleConfig(cfg *sc.ScheduleConfig) {
 	}
 }
 
+// UpdateScheduleConfig updates the latest scheduling configuration.
+func (o *PersistConfig) UpdateScheduleConfig(storage endpoint.ConfigStorage, update func(*sc.ScheduleConfig) error) error {
+	o.configMu.Lock()
+	defer o.configMu.Unlock()
+
+	old := o.GetScheduleConfig()
+	cfg := old.Clone()
+	if err := update(cfg); err != nil {
+		return err
+	}
+	o.storeScheduleConfig(cfg)
+	if err := o.Persist(storage); err != nil {
+		o.storeScheduleConfig(old)
+		return err
+	}
+	return nil
+}
+
 // AdjustScheduleCfg adjusts the schedule config during the initialization.
 func AdjustScheduleCfg(scheduleCfg *sc.ScheduleConfig) {
 	// In case we add new default schedulers.

pkg/mcs/scheduling/server/server.go

@@ -190,47 +190,57 @@ func (s *Server) updatePDMemberLoop() {
 		case <-ticker.C:
 		case <-s.checkMembershipCh:
 		}
-		if !s.IsServing() {
+		curLeader = s.updatePDLeader(ctx, curLeader)
+	}
+}
+
+func (s *Server) updatePDLeader(ctx context.Context, curLeader uint64) uint64 {
+	if !s.IsServing() {
+		return curLeader
+	}
+	members, err := etcdutil.ListEtcdMembers(ctx, s.GetClient())
+	if err != nil {
+		log.Warn("failed to list members", errs.ZapError(err))
+		return curLeader
+	}
+	for _, ep := range members.Members {
+		if len(ep.GetClientURLs()) == 0 { // This member is not started yet.
+			log.Info("member is not started yet", zap.String("member-id", strconv.FormatUint(ep.GetID(), 16)), errs.ZapError(err))
 			continue
 		}
-		members, err := etcdutil.ListEtcdMembers(ctx, s.GetClient())
+		probeCtx, cancel := context.WithTimeout(ctx, requestTimeout)
+		status, err := s.GetClient().Status(probeCtx, ep.ClientURLs[0])
+		cancel()
 		if err != nil {
-			log.Warn("failed to list members", errs.ZapError(err))
+			log.Info("failed to get status of member", zap.String("member-id", strconv.FormatUint(ep.ID, 16)), zap.String("endpoint", ep.ClientURLs[0]), errs.ZapError(err))
 			continue
 		}
-		for _, ep := range members.Members {
-			if len(ep.GetClientURLs()) == 0 { // This member is not started yet.
-				log.Info("member is not started yet", zap.String("member-id", strconv.FormatUint(ep.GetID(), 16)), errs.ZapError(err))
-				continue
-			}
-			status, err := s.GetClient().Status(ctx, ep.ClientURLs[0])
-			if err != nil {
-				log.Info("failed to get status of member", zap.String("member-id", strconv.FormatUint(ep.ID, 16)), zap.String("endpoint", ep.ClientURLs[0]), errs.ZapError(err))
-				continue
-			}
-			if status.Leader == ep.ID {
-				cc, err := s.GetDelegateClient(ctx, s.GetTLSConfig(), ep.ClientURLs[0])
-				if err != nil {
-					log.Info("failed to get delegate client", errs.ZapError(err))
-					continue
-				}
-				if !s.IsServing() {
-					// double check
-					break
-				}
-				cluster := s.GetCluster()
-				if cluster != nil {
-					if cluster.SwitchPDLeader(pdpb.NewPDClient(cc)) {
-						if status.Leader != curLeader {
-							log.Info("switch PD leader", zap.String("leader-id", strconv.FormatUint(ep.ID, 16)), zap.String("endpoint", ep.ClientURLs[0]))
-						}
-						curLeader = ep.ID
-						break
-					}
-				}
+		if status.Leader != ep.ID {
+			continue
+		}
+		probeCtx, cancel = context.WithTimeout(ctx, requestTimeout)
+		cc, err := s.GetDelegateClient(probeCtx, s.GetTLSConfig(), ep.ClientURLs[0])
+		cancel()
+		if err != nil {
+			log.Info("failed to get delegate client", errs.ZapError(err))
+			continue
+		}
+		if !s.IsServing() {
+			// double check
+			return curLeader
+		}
+		cluster := s.GetCluster()
+		if cluster == nil {
+			return curLeader
+		}
+		if cluster.SwitchPDLeader(pdpb.NewPDClient(cc)) {
+			if status.Leader != curLeader {
+				log.Info("switch PD leader", zap.String("leader-id", strconv.FormatUint(ep.ID, 16)), zap.String("endpoint", ep.ClientURLs[0]))
 			}
+			return ep.ID
 		}
 	}
+	return curLeader
 }
 
 func (s *Server) primaryElectionLoop() {

pkg/schedule/config/config_provider.go

@@ -69,6 +69,7 @@ type SchedulerConfigProvider interface {
 
 	GetScheduleConfig() *ScheduleConfig
 	SetScheduleConfig(*ScheduleConfig)
+	UpdateScheduleConfig(endpoint.ConfigStorage, func(*ScheduleConfig) error) error
 }
 
 // CheckerConfigProvider is the interface for checker configurations.

pkg/schedule/coordinator.go

@@ -16,6 +16,7 @@ package schedule
 
 import (
 	"context"
+	"reflect"
 	"strconv"
 	"sync"
 	"time"
@@ -38,6 +39,7 @@ import (
 	"github.com/tikv/pd/pkg/schedule/schedulers"
 	"github.com/tikv/pd/pkg/schedule/splitter"
 	"github.com/tikv/pd/pkg/schedule/types"
+	"github.com/tikv/pd/pkg/slice"
 	"github.com/tikv/pd/pkg/statistics"
 	"github.com/tikv/pd/pkg/statistics/utils"
 	"github.com/tikv/pd/pkg/utils/logutil"
@@ -283,6 +285,7 @@ func (c *Coordinator) InitSchedulers(needRun bool) {
 		log.Fatal("cannot load schedulers' config", errs.ZapError(err))
 	}
 	scheduleCfg := c.cluster.GetSchedulerConfig().GetScheduleConfig().Clone()
+	originalSchedulerCfgs := append(sc.SchedulerConfigs(nil), scheduleCfg.Schedulers...)
 	// The new way to create scheduler with the independent configuration.
 	for i, name := range scheduleNames {
 		select {
@@ -318,13 +321,13 @@ func (c *Coordinator) InitSchedulers(needRun bool) {
 		}
 		if needRun {
 			log.Info("create scheduler with independent configuration", zap.String("scheduler-name", s.GetName()))
-			if err = c.schedulers.AddScheduler(s); err != nil {
+			if err = c.schedulers.AddScheduler(s, cfg.Args...); 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 {
 			log.Info("create scheduler handler with independent configuration", zap.String("scheduler-name", s.GetName()))
-			if err = c.schedulers.AddSchedulerHandler(s); err != nil {
+			if err = c.schedulers.AddSchedulerHandler(s, cfg.Args...); 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))
 			}
@@ -382,16 +385,46 @@ func (c *Coordinator) InitSchedulers(needRun bool) {
 	}
 
 	// Removes the invalid scheduler config and persist.
-	scheduleCfg.Schedulers = scheduleCfg.Schedulers[:k]
-	c.cluster.GetSchedulerConfig().SetScheduleConfig(scheduleCfg)
-	if err := c.cluster.GetSchedulerConfig().Persist(c.cluster.GetStorage()); err != nil {
+	validSchedulerCfgs := scheduleCfg.Schedulers[:k]
+	invalidSchedulerCfgs := make(sc.SchedulerConfigs, 0, len(originalSchedulerCfgs)-len(validSchedulerCfgs))
+	for _, schedulerCfg := range originalSchedulerCfgs {
+		if !containsSchedulerConfig(validSchedulerCfgs, schedulerCfg) {
+			invalidSchedulerCfgs = append(invalidSchedulerCfgs, schedulerCfg)
+		}
+	}
+	var updatedSchedulers sc.SchedulerConfigs
+	if err := c.cluster.GetSchedulerConfig().UpdateScheduleConfig(c.cluster.GetStorage(), func(latest *sc.ScheduleConfig) error {
+		if len(invalidSchedulerCfgs) == 0 {
+			updatedSchedulers = append(sc.SchedulerConfigs(nil), latest.Schedulers...)
+			return nil
+		}
+		latest.Schedulers = removeSchedulerConfigs(latest.Schedulers, invalidSchedulerCfgs)
+		updatedSchedulers = append(sc.SchedulerConfigs(nil), latest.Schedulers...)
+		return nil
+	}); err != nil {
 		log.Error("cannot persist schedule config", errs.ZapError(err))
 	}
-	log.Info("scheduler config is updated", zap.Reflect("scheduler-config", scheduleCfg.Schedulers))
+	log.Info("scheduler config is updated", zap.Reflect("scheduler-config", updatedSchedulers))
 
 	c.markSchedulersInitialized()
 }
 
+func containsSchedulerConfig(schedulers sc.SchedulerConfigs, target sc.SchedulerConfig) bool {
+	return slice.AnyOf(schedulers, func(i int) bool {
+		return reflect.DeepEqual(schedulers[i], target)
+	})
+}
+
+func removeSchedulerConfigs(schedulers, removed sc.SchedulerConfigs) sc.SchedulerConfigs {
+	retained := schedulers[:0]
+	for _, schedulerCfg := range schedulers {
+		if !containsSchedulerConfig(removed, schedulerCfg) {
+			retained = append(retained, schedulerCfg)
+		}
+	}
+	return retained
+}
+
 // LoadPlugin load user plugin
 func (c *Coordinator) LoadPlugin(pluginPath string, ch chan string) {
 	log.Info("load plugin", zap.String("plugin-path", pluginPath))

pkg/schedule/coordinator_test.go

@@ -0,0 +1,58 @@
+// Copyright 2026 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package schedule
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
+
+	sc "github.com/tikv/pd/pkg/schedule/config"
+	"github.com/tikv/pd/pkg/schedule/types"
+	"github.com/tikv/pd/pkg/utils/testutil"
+)
+
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
+func TestRemoveSchedulerConfigsKeepsConcurrentUpdates(t *testing.T) {
+	re := require.New(t)
+
+	invalidOldCfg := sc.SchedulerConfig{
+		Type: types.SchedulerTypeCompatibleMap[types.GrantLeaderScheduler],
+		Args: []string{"1"},
+	}
+	concurrentUpdatedCfg := sc.SchedulerConfig{
+		Type: types.SchedulerTypeCompatibleMap[types.GrantLeaderScheduler],
+		Args: []string{"2"},
+	}
+	concurrentAddedCfg := sc.SchedulerConfig{
+		Type: types.SchedulerTypeCompatibleMap[types.ShuffleRegionScheduler],
+	}
+	latest := sc.SchedulerConfigs{
+		sc.DefaultSchedulers[0],
+		invalidOldCfg,
+		concurrentUpdatedCfg,
+		concurrentAddedCfg,
+	}
+
+	re.Equal(sc.SchedulerConfigs{
+		sc.DefaultSchedulers[0],
+		concurrentUpdatedCfg,
+		concurrentAddedCfg,
+	}, removeSchedulerConfigs(latest, sc.SchedulerConfigs{invalidOldCfg}))
+}