From 773fb3784589a2c4f8657a0001bb6b09211649b3 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Wed, 10 Jun 2026 14:46:24 +0800 Subject: [PATCH 1/2] mcs: prevent scheduler update notifier from blocking Signed-off-by: Ryan Leung --- pkg/mcs/scheduling/server/cluster.go | 1 + pkg/mcs/scheduling/server/config/config.go | 28 +++++-- .../scheduling/server/config/config_test.go | 84 +++++++++++++++++++ 3 files changed, 108 insertions(+), 5 deletions(-) create mode 100644 pkg/mcs/scheduling/server/config/config_test.go diff --git a/pkg/mcs/scheduling/server/cluster.go b/pkg/mcs/scheduling/server/cluster.go index 0ae112e18a2..fdfa1c5e347 100644 --- a/pkg/mcs/scheduling/server/cluster.go +++ b/pkg/mcs/scheduling/server/cluster.go @@ -425,6 +425,7 @@ func (c *Cluster) updateScheduler() { // Make sure the check will be triggered once later. trySend(notifier) c.persistConfig.SetSchedulersUpdatingNotifier(notifier) + defer c.persistConfig.ClearSchedulersUpdatingNotifier(notifier) ticker := time.NewTicker(time.Second) defer ticker.Stop() diff --git a/pkg/mcs/scheduling/server/config/config.go b/pkg/mcs/scheduling/server/config/config.go index e04bea49e2e..f7aeb6b4bd0 100644 --- a/pkg/mcs/scheduling/server/config/config.go +++ b/pkg/mcs/scheduling/server/config/config.go @@ -197,7 +197,7 @@ type PersistConfig struct { replication atomic.Value storeConfig atomic.Value // schedulersUpdatingNotifier is used to notify that the schedulers have been updated. - // Store as `chan<- struct{}`. + // Store as `chan struct{}`. schedulersUpdatingNotifier atomic.Value } @@ -215,16 +215,30 @@ func NewPersistConfig(cfg *Config, ttl *cache.TTLString) *PersistConfig { } // SetSchedulersUpdatingNotifier sets the schedulers updating notifier. -func (o *PersistConfig) SetSchedulersUpdatingNotifier(notifier chan<- struct{}) { +func (o *PersistConfig) SetSchedulersUpdatingNotifier(notifier chan struct{}) { o.schedulersUpdatingNotifier.Store(notifier) } -func (o *PersistConfig) getSchedulersUpdatingNotifier() chan<- struct{} { +// ClearSchedulersUpdatingNotifier clears the schedulers updating notifier if it +// still matches the given notifier. +func (o *PersistConfig) ClearSchedulersUpdatingNotifier(notifier chan struct{}) { + if notifier == nil { + return + } + current := o.getSchedulersUpdatingNotifier() + if current != notifier { + return + } + var empty chan struct{} + o.schedulersUpdatingNotifier.CompareAndSwap(current, empty) +} + +func (o *PersistConfig) getSchedulersUpdatingNotifier() chan struct{} { v := o.schedulersUpdatingNotifier.Load() if v == nil { return nil } - return v.(chan<- struct{}) + return v.(chan struct{}) } func (o *PersistConfig) tryNotifySchedulersUpdating() { @@ -232,7 +246,11 @@ func (o *PersistConfig) tryNotifySchedulersUpdating() { if notifier == nil { return } - notifier <- struct{}{} + // Scheduler update notifications are coalesced; one pending signal is enough. + select { + case notifier <- struct{}{}: + default: + } } // GetClusterVersion returns the cluster version. diff --git a/pkg/mcs/scheduling/server/config/config_test.go b/pkg/mcs/scheduling/server/config/config_test.go new file mode 100644 index 00000000000..d566ce7249c --- /dev/null +++ b/pkg/mcs/scheduling/server/config/config_test.go @@ -0,0 +1,84 @@ +// 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 config + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" + + sc "github.com/tikv/pd/pkg/schedule/config" + "github.com/tikv/pd/pkg/schedule/types" +) + +func TestSetScheduleConfigDoesNotBlockWhenSchedulersUpdatingNotifierIsFull(t *testing.T) { + re := require.New(t) + persistConfig := NewPersistConfig(&Config{}, nil) + notifier := make(chan struct{}, 1) + notifier <- struct{}{} + persistConfig.SetSchedulersUpdatingNotifier(notifier) + + done := make(chan struct{}) + go func() { + persistConfig.SetScheduleConfig(newScheduleConfig(types.BalanceLeaderScheduler)) + close(done) + }() + + select { + case <-done: + case <-time.After(time.Second): + re.Fail("SetScheduleConfig should not block on a full schedulers updating notifier") + } +} + +func TestClearSchedulersUpdatingNotifier(t *testing.T) { + re := require.New(t) + persistConfig := NewPersistConfig(&Config{}, nil) + oldNotifier := make(chan struct{}, 1) + newNotifier := make(chan struct{}, 1) + + persistConfig.SetSchedulersUpdatingNotifier(oldNotifier) + persistConfig.SetSchedulersUpdatingNotifier(newNotifier) + persistConfig.ClearSchedulersUpdatingNotifier(oldNotifier) + persistConfig.SetScheduleConfig(newScheduleConfig(types.BalanceLeaderScheduler)) + + select { + case <-oldNotifier: + re.Fail("old notifier should not receive scheduler update after a newer notifier is registered") + default: + } + select { + case <-newNotifier: + case <-time.After(time.Second): + re.Fail("new notifier should receive scheduler update") + } + + persistConfig.ClearSchedulersUpdatingNotifier(newNotifier) + persistConfig.SetScheduleConfig(newScheduleConfig(types.BalanceRegionScheduler)) + select { + case <-newNotifier: + re.Fail("cleared notifier should not receive scheduler update") + default: + } +} + +func newScheduleConfig(schedulerType types.CheckerSchedulerType) *sc.ScheduleConfig { + return &sc.ScheduleConfig{ + Schedulers: []sc.SchedulerConfig{ + {Type: types.SchedulerTypeCompatibleMap[schedulerType]}, + }, + } +} From e164654345b6a6160a484f2c12561a1a203bbf23 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Wed, 10 Jun 2026 15:06:29 +0800 Subject: [PATCH 2/2] mcs: cover scheduler notifier tests with goleak Signed-off-by: Ryan Leung --- pkg/mcs/scheduling/server/config/config_test.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pkg/mcs/scheduling/server/config/config_test.go b/pkg/mcs/scheduling/server/config/config_test.go index d566ce7249c..9353bce3269 100644 --- a/pkg/mcs/scheduling/server/config/config_test.go +++ b/pkg/mcs/scheduling/server/config/config_test.go @@ -19,11 +19,17 @@ import ( "time" "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 TestSetScheduleConfigDoesNotBlockWhenSchedulersUpdatingNotifierIsFull(t *testing.T) { re := require.New(t) persistConfig := NewPersistConfig(&Config{}, nil)