From db0ecac4ae5e5fff5977f0e47de09e586c35e58c Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Sat, 7 Mar 2026 19:49:36 +0800 Subject: [PATCH 1/2] ai scan potential panix Signed-off-by: tongjian <1045931706@qq.com> --- pkg/schedule/schedulers/balance_range.go | 46 ++++++++++---- pkg/schedule/schedulers/balance_range_test.go | 41 +++++++++++++ pkg/schedule/schedulers/evict_leader.go | 35 ++++++++--- pkg/schedule/schedulers/evict_leader_test.go | 59 ++++++++++++++++++ pkg/schedule/schedulers/grant_hot_region.go | 7 ++- .../schedulers/grant_hot_region_test.go | 45 ++++++++++++++ pkg/schedule/schedulers/grant_leader.go | 28 ++++++++- pkg/schedule/schedulers/grant_leader_test.go | 60 +++++++++++++++++++ .../schedulers/scheduler_controller.go | 7 ++- .../schedulers/transfer_witness_leader.go | 5 +- .../transfer_witness_leader_test.go | 8 +++ server/api/scheduler.go | 26 +++++--- 12 files changed, 337 insertions(+), 30 deletions(-) create mode 100644 pkg/schedule/schedulers/grant_hot_region_test.go diff --git a/pkg/schedule/schedulers/balance_range.go b/pkg/schedule/schedulers/balance_range.go index b80d6109ec2..ff810ca9db6 100644 --- a/pkg/schedule/schedulers/balance_range.go +++ b/pkg/schedule/schedulers/balance_range.go @@ -90,27 +90,48 @@ func (handler *balanceRangeSchedulerHandler) addJob(w http.ResponseWriter, r *ht Status: pending, Timeout: defaultJobTimeout, } - job.Engine = input["engine"].(string) + engine, ok := input["engine"].(string) + if !ok || len(engine) == 0 { + handler.rd.JSON(w, http.StatusBadRequest, "engine is required and must be a string") + return + } + job.Engine = engine if job.Engine != core.EngineTiFlash && job.Engine != core.EngineTiKV { - handler.rd.JSON(w, http.StatusBadRequest, fmt.Sprintf("engine:%s must be tikv or tiflash", input["engine"].(string))) + handler.rd.JSON(w, http.StatusBadRequest, fmt.Sprintf("engine:%s must be tikv or tiflash", job.Engine)) + return + } + ruleStr, ok := input["rule"].(string) + if !ok || len(ruleStr) == 0 { + handler.rd.JSON(w, http.StatusBadRequest, "rule is required and must be a string") return } - job.Rule = core.NewRule(input["rule"].(string)) + job.Rule = core.NewRule(ruleStr) if job.Rule != core.LeaderScatter && job.Rule != core.PeerScatter && job.Rule != core.LearnerScatter { handler.rd.JSON(w, http.StatusBadRequest, fmt.Sprintf("rule:%s must be leader-scatter, learner-scatter or peer-scatter", - input["engine"].(string))) + ruleStr)) return } - job.Alias = input["alias"].(string) - timeoutStr, ok := input["timeout"].(string) - if ok && len(timeoutStr) > 0 { - timeout, err := time.ParseDuration(timeoutStr) - if err != nil { - handler.rd.JSON(w, http.StatusBadRequest, fmt.Sprintf("timeout:%s is invalid", input["timeout"].(string))) + alias, ok := input["alias"].(string) + if !ok || len(alias) == 0 { + handler.rd.JSON(w, http.StatusBadRequest, "alias is required and must be a string") + return + } + job.Alias = alias + if timeoutVal, exists := input["timeout"]; exists { + timeoutStr, ok := timeoutVal.(string) + if !ok { + handler.rd.JSON(w, http.StatusBadRequest, "timeout must be a string") return } - job.Timeout = timeout + if len(timeoutStr) > 0 { + timeout, err := time.ParseDuration(timeoutStr) + if err != nil { + handler.rd.JSON(w, http.StatusBadRequest, fmt.Sprintf("timeout:%s is invalid", timeoutStr)) + return + } + job.Timeout = timeout + } } keys, err := keyutil.DecodeHTTPKeyRanges(input) @@ -327,6 +348,9 @@ func (job *balanceRangeSchedulerJob) expired(dur time.Duration) bool { } func (job *balanceRangeSchedulerJob) shouldFinished() bool { + if job == nil || job.Start == nil { + return true + } return time.Since(*job.Start) > job.Timeout } diff --git a/pkg/schedule/schedulers/balance_range_test.go b/pkg/schedule/schedulers/balance_range_test.go index 7b729ae8903..bea3b6eaf1e 100644 --- a/pkg/schedule/schedulers/balance_range_test.go +++ b/pkg/schedule/schedulers/balance_range_test.go @@ -15,7 +15,11 @@ package schedulers import ( + "bytes" + "encoding/json" "fmt" + "net/http" + "net/http/httptest" "strconv" "testing" "time" @@ -23,6 +27,7 @@ import ( "github.com/stretchr/testify/require" "github.com/pingcap/failpoint" + "github.com/unrolled/render" "github.com/tikv/pd/pkg/core" "github.com/tikv/pd/pkg/schedule/operator" @@ -516,3 +521,39 @@ func TestPersistFail(t *testing.T) { re.ErrorContains(conf.gcLocked(), errMsg) re.Len(conf.jobs, 1) } + +func TestAddBalanceRangeJobWithInvalidFieldType(t *testing.T) { + re := require.New(t) + conf := &balanceRangeSchedulerConfig{ + schedulerConfig: &baseSchedulerConfig{}, + jobs: make([]*balanceRangeSchedulerJob, 0), + } + conf.init("test", storage.NewStorageWithMemoryBackend(), conf) + handler := &balanceRangeSchedulerHandler{ + config: conf, + rd: render.New(render.Options{IndentJSON: true}), + } + body, err := json.Marshal(map[string]any{ + "alias": "a", + "engine": 1, + "rule": "leader-scatter", + "start-key": "100", + "end-key": "200", + }) + re.NoError(err) + req := httptest.NewRequest(http.MethodPut, "/job", bytes.NewReader(body)) + resp := httptest.NewRecorder() + re.NotPanics(func() { + handler.addJob(resp, req) + }) + re.Equal(http.StatusBadRequest, resp.Code) + re.Empty(conf.jobs) +} + +func TestJobShouldFinishedWhenStartTimeIsMissing(t *testing.T) { + re := require.New(t) + job := &balanceRangeSchedulerJob{Timeout: time.Minute} + re.True(job.shouldFinished()) + job = nil + re.True(job.shouldFinished()) +} diff --git a/pkg/schedule/schedulers/evict_leader.go b/pkg/schedule/schedulers/evict_leader.go index 7f825bdba00..5adcbee79e9 100644 --- a/pkg/schedule/schedulers/evict_leader.go +++ b/pkg/schedule/schedulers/evict_leader.go @@ -173,12 +173,13 @@ func (conf *evictLeaderSchedulerConfig) resumeLeaderTransfer(cluster sche.Schedu func (conf *evictLeaderSchedulerConfig) pauseLeaderTransferIfStoreNotExist(id uint64) (bool, error) { conf.RLock() defer conf.RUnlock() - if _, exist := conf.StoreIDWithRanges[id]; !exist { - if err := conf.cluster.PauseLeaderTransfer(id, constant.In); err != nil { - return exist, err - } + if _, exist := conf.StoreIDWithRanges[id]; exist { + return true, nil + } + if err := conf.cluster.PauseLeaderTransfer(id, constant.In); err != nil { + return false, err } - return true, nil + return false, nil } func (conf *evictLeaderSchedulerConfig) resumeLeaderTransferIfExist(id uint64) { @@ -427,8 +428,28 @@ func (handler *evictLeaderHandler) updateConfig(w http.ResponseWriter, r *http.R batch = (int)(batchFloat) } - ranges, ok := (input["ranges"]).([]string) - if ok { + var ranges []string + rangesVal, hasRanges := input["ranges"] + if hasRanges { + switch val := rangesVal.(type) { + case []string: + ranges = val + case []any: + ranges = make([]string, 0, len(val)) + for _, item := range val { + s, ok := item.(string) + if !ok { + handler.config.resumeLeaderTransferIfExist(id) + handler.rd.JSON(w, http.StatusBadRequest, errs.ErrSchedulerConfig.FastGenByArgs("ranges")) + return + } + ranges = append(ranges, s) + } + default: + handler.config.resumeLeaderTransferIfExist(id) + handler.rd.JSON(w, http.StatusBadRequest, errs.ErrSchedulerConfig.FastGenByArgs("ranges")) + return + } if !inputHasStoreID { handler.config.resumeLeaderTransferIfExist(id) handler.rd.JSON(w, http.StatusBadRequest, errs.ErrSchedulerConfig.FastGenByArgs("id")) diff --git a/pkg/schedule/schedulers/evict_leader_test.go b/pkg/schedule/schedulers/evict_leader_test.go index c665a548062..c47d6121172 100644 --- a/pkg/schedule/schedulers/evict_leader_test.go +++ b/pkg/schedule/schedulers/evict_leader_test.go @@ -17,9 +17,12 @@ package schedulers import ( "bytes" "encoding/json" + "net/http" + "net/http/httptest" "testing" "github.com/stretchr/testify/require" + "github.com/unrolled/render" "github.com/pingcap/errors" "github.com/pingcap/failpoint" @@ -251,3 +254,59 @@ func TestEvictLeaderDeleteWithSaveFailure(t *testing.T) { re.Equal(keyRanges, conf.StoreIDWithRanges[1], "key ranges should be restored") re.Empty(resp) } + +func TestEvictLeaderUpdateConfigWithStringArrayRanges(t *testing.T) { + re := require.New(t) + cancel, _, tc, _ := prepareSchedulersTest() + defer cancel() + + tc.AddLeaderStore(1, 0) + conf := &evictLeaderSchedulerConfig{ + schedulerConfig: &baseSchedulerConfig{}, + StoreIDWithRanges: make(map[uint64][]keyutil.KeyRange), + Batch: EvictLeaderBatchSize, + cluster: tc.GetBasicCluster(), + } + conf.init("evict-leader-test", storage.NewStorageWithMemoryBackend(), conf) + handler := &evictLeaderHandler{config: conf, rd: render.New(render.Options{IndentJSON: true})} + body, err := json.Marshal(map[string]any{ + "store_id": 1, + "ranges": []string{"100", "200"}, + }) + re.NoError(err) + req := httptest.NewRequest(http.MethodPost, "/config", bytes.NewReader(body)) + resp := httptest.NewRecorder() + re.NotPanics(func() { + handler.updateConfig(resp, req) + }) + re.Equal(http.StatusOK, resp.Code) + re.Len(conf.StoreIDWithRanges[1], 1) + re.Equal(keyutil.NewKeyRange("100", "200"), conf.StoreIDWithRanges[1][0]) +} + +func TestEvictLeaderUpdateConfigWithInvalidRangesType(t *testing.T) { + re := require.New(t) + cancel, _, tc, _ := prepareSchedulersTest() + defer cancel() + + tc.AddLeaderStore(1, 0) + conf := &evictLeaderSchedulerConfig{ + schedulerConfig: &baseSchedulerConfig{}, + StoreIDWithRanges: make(map[uint64][]keyutil.KeyRange), + Batch: EvictLeaderBatchSize, + cluster: tc.GetBasicCluster(), + } + conf.init("evict-leader-test", storage.NewStorageWithMemoryBackend(), conf) + handler := &evictLeaderHandler{config: conf, rd: render.New(render.Options{IndentJSON: true})} + body, err := json.Marshal(map[string]any{ + "store_id": 1, + "ranges": []any{"100", 200}, + }) + re.NoError(err) + req := httptest.NewRequest(http.MethodPost, "/config", bytes.NewReader(body)) + resp := httptest.NewRecorder() + re.NotPanics(func() { + handler.updateConfig(resp, req) + }) + re.Equal(http.StatusBadRequest, resp.Code) +} diff --git a/pkg/schedule/schedulers/grant_hot_region.go b/pkg/schedule/schedulers/grant_hot_region.go index 7c169be134b..18a70de7926 100644 --- a/pkg/schedule/schedulers/grant_hot_region.go +++ b/pkg/schedule/schedulers/grant_hot_region.go @@ -189,7 +189,12 @@ func (handler *grantHotRegionHandler) updateConfig(w http.ResponseWriter, r *htt } storeIDs = append(storeIDs, id) } - leaderID, err := strconv.ParseUint(input["store-leader-id"].(string), 10, 64) + leaderStr, ok := input["store-leader-id"].(string) + if !ok { + handler.rd.JSON(w, http.StatusBadRequest, errs.ErrSchedulerConfig) + return + } + leaderID, err := strconv.ParseUint(leaderStr, 10, 64) if err != nil { handler.rd.JSON(w, http.StatusBadRequest, errs.ErrBytesToUint64) return diff --git a/pkg/schedule/schedulers/grant_hot_region_test.go b/pkg/schedule/schedulers/grant_hot_region_test.go new file mode 100644 index 00000000000..402dec78e80 --- /dev/null +++ b/pkg/schedule/schedulers/grant_hot_region_test.go @@ -0,0 +1,45 @@ +// 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 schedulers + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" + "github.com/unrolled/render" +) + +func TestGrantHotRegionUpdateConfigWithInvalidLeaderIDType(t *testing.T) { + re := require.New(t) + handler := &grantHotRegionHandler{ + config: &grantHotRegionSchedulerConfig{}, + rd: render.New(render.Options{IndentJSON: true}), + } + body, err := json.Marshal(map[string]any{ + "store-id": "1,2", + "store-leader-id": 1, + }) + re.NoError(err) + req := httptest.NewRequest(http.MethodPost, "/config", bytes.NewReader(body)) + resp := httptest.NewRecorder() + re.NotPanics(func() { + handler.updateConfig(resp, req) + }) + re.Equal(http.StatusBadRequest, resp.Code) +} diff --git a/pkg/schedule/schedulers/grant_leader.go b/pkg/schedule/schedulers/grant_leader.go index abd18304db9..97c02285c2c 100644 --- a/pkg/schedule/schedulers/grant_leader.go +++ b/pkg/schedule/schedulers/grant_leader.go @@ -265,8 +265,32 @@ func (handler *grantLeaderHandler) updateConfig(w http.ResponseWriter, r *http.R args = append(args, strconv.FormatUint(id, 10)) } - ranges, ok := (input["ranges"]).([]string) - if ok { + rangesVal, hasRanges := input["ranges"] + if hasRanges { + var ranges []string + switch val := rangesVal.(type) { + case []string: + ranges = val + case []any: + ranges = make([]string, 0, len(val)) + for _, item := range val { + s, ok := item.(string) + if !ok { + handler.config.Lock() + handler.config.cluster.ResumeLeaderTransfer(id, constant.Out) + handler.config.Unlock() + handler.rd.JSON(w, http.StatusBadRequest, errs.ErrSchedulerConfig.FastGenByArgs("ranges")) + return + } + ranges = append(ranges, s) + } + default: + handler.config.Lock() + handler.config.cluster.ResumeLeaderTransfer(id, constant.Out) + handler.config.Unlock() + handler.rd.JSON(w, http.StatusBadRequest, errs.ErrSchedulerConfig.FastGenByArgs("ranges")) + return + } args = append(args, ranges...) } else if exists { args = append(args, handler.config.getRanges(id)...) diff --git a/pkg/schedule/schedulers/grant_leader_test.go b/pkg/schedule/schedulers/grant_leader_test.go index 003966db73d..34718c4bedd 100644 --- a/pkg/schedule/schedulers/grant_leader_test.go +++ b/pkg/schedule/schedulers/grant_leader_test.go @@ -15,13 +15,19 @@ package schedulers import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" "testing" "github.com/stretchr/testify/require" + "github.com/unrolled/render" "github.com/tikv/pd/pkg/schedule/operator" "github.com/tikv/pd/pkg/schedule/types" "github.com/tikv/pd/pkg/storage" + "github.com/tikv/pd/pkg/utils/keyutil" "github.com/tikv/pd/pkg/utils/operatorutil" ) @@ -60,3 +66,57 @@ func TestGrantLeaderScheduler(t *testing.T) { ops, _ = bls.Schedule(tc, false) re.Empty(ops) } + +func TestGrantLeaderUpdateConfigWithStringArrayRanges(t *testing.T) { + re := require.New(t) + cancel, _, tc, _ := prepareSchedulersTest() + defer cancel() + + tc.AddLeaderStore(1, 0) + conf := &grantLeaderSchedulerConfig{ + schedulerConfig: &baseSchedulerConfig{}, + StoreIDWithRanges: make(map[uint64][]keyutil.KeyRange), + cluster: tc.GetBasicCluster(), + } + conf.init("grant-leader-test", storage.NewStorageWithMemoryBackend(), conf) + handler := &grantLeaderHandler{config: conf, rd: render.New(render.Options{IndentJSON: true})} + body, err := json.Marshal(map[string]any{ + "store_id": 1, + "ranges": []string{"100", "200"}, + }) + re.NoError(err) + req := httptest.NewRequest(http.MethodPost, "/config", bytes.NewReader(body)) + resp := httptest.NewRecorder() + re.NotPanics(func() { + handler.updateConfig(resp, req) + }) + re.Equal(http.StatusOK, resp.Code) + re.Len(conf.StoreIDWithRanges[1], 1) + re.Equal(keyutil.NewKeyRange("100", "200"), conf.StoreIDWithRanges[1][0]) +} + +func TestGrantLeaderUpdateConfigWithInvalidRangesType(t *testing.T) { + re := require.New(t) + cancel, _, tc, _ := prepareSchedulersTest() + defer cancel() + + tc.AddLeaderStore(1, 0) + conf := &grantLeaderSchedulerConfig{ + schedulerConfig: &baseSchedulerConfig{}, + StoreIDWithRanges: make(map[uint64][]keyutil.KeyRange), + cluster: tc.GetBasicCluster(), + } + conf.init("grant-leader-test", storage.NewStorageWithMemoryBackend(), conf) + handler := &grantLeaderHandler{config: conf, rd: render.New(render.Options{IndentJSON: true})} + body, err := json.Marshal(map[string]any{ + "store_id": 1, + "ranges": []any{"100", 200}, + }) + re.NoError(err) + req := httptest.NewRequest(http.MethodPost, "/config", bytes.NewReader(body)) + resp := httptest.NewRecorder() + re.NotPanics(func() { + handler.updateConfig(resp, req) + }) + re.Equal(http.StatusBadRequest, resp.Code) +} diff --git a/pkg/schedule/schedulers/scheduler_controller.go b/pkg/schedule/schedulers/scheduler_controller.go index cbfdddddaa0..faf67cdd5a7 100644 --- a/pkg/schedule/schedulers/scheduler_controller.go +++ b/pkg/schedule/schedulers/scheduler_controller.go @@ -429,8 +429,13 @@ func (c *Controller) CheckTransferWitnessLeader(region *core.RegionInfo) { s, ok := c.schedulers[types.TransferWitnessLeaderScheduler.String()] c.RUnlock() if ok { + regionC := RecvRegionInfo(s.Scheduler) + if regionC == nil { + log.Warn("invalid scheduler type for transfer witness leader", zap.String("scheduler", s.GetName())) + return + } select { - case RecvRegionInfo(s.Scheduler) <- region: + case regionC <- region: default: log.Warn("drop transfer witness leader due to recv region channel full", zap.Uint64("region-id", region.GetID())) } diff --git a/pkg/schedule/schedulers/transfer_witness_leader.go b/pkg/schedule/schedulers/transfer_witness_leader.go index be583c40ae2..1289de36833 100644 --- a/pkg/schedule/schedulers/transfer_witness_leader.go +++ b/pkg/schedule/schedulers/transfer_witness_leader.go @@ -114,5 +114,8 @@ func scheduleTransferWitnessLeader(name string, cluster sche.SchedulerCluster, r // RecvRegionInfo receives a checked region from coordinator func RecvRegionInfo(s Scheduler) chan<- *core.RegionInfo { - return s.(*transferWitnessLeaderScheduler).regions + if scheduler, ok := s.(*transferWitnessLeaderScheduler); ok { + return scheduler.regions + } + return nil } diff --git a/pkg/schedule/schedulers/transfer_witness_leader_test.go b/pkg/schedule/schedulers/transfer_witness_leader_test.go index 14d40f33cfc..0ffdc3c9fa5 100644 --- a/pkg/schedule/schedulers/transfer_witness_leader_test.go +++ b/pkg/schedule/schedulers/transfer_witness_leader_test.go @@ -89,3 +89,11 @@ func TestTransferWitnessLeaderWithUnhealthyPeer(t *testing.T) { ops, _ = sl.Schedule(tc, false) re.Empty(ops) } + +func TestRecvRegionInfoWithWrongSchedulerType(t *testing.T) { + re := require.New(t) + var s Scheduler = &balanceRangeScheduler{} + re.NotPanics(func() { + re.Nil(RecvRegionInfo(s)) + }) +} diff --git a/server/api/scheduler.go b/server/api/scheduler.go index 6610c448aaa..b4a51d47ef8 100644 --- a/server/api/scheduler.go +++ b/server/api/scheduler.go @@ -125,15 +125,19 @@ func (h *schedulerHandler) CreateScheduler(w http.ResponseWriter, r *http.Reques handler.ServeHTTP(w, r) return } - h.r.JSON(w, http.StatusNotAcceptable, err.Error()) + if err != nil { + h.r.JSON(w, http.StatusNotAcceptable, err.Error()) + return + } + h.r.JSON(w, http.StatusNotAcceptable, "scheduler config handler is unavailable") return } if err := apiutil.CollectStringOption("rule", input, collector); err != nil { - h.r.JSON(w, http.StatusInternalServerError, err.Error()) + h.r.JSON(w, http.StatusBadRequest, err.Error()) return } if err := apiutil.CollectStringOption("engine", input, collector); err != nil { - h.r.JSON(w, http.StatusInternalServerError, err.Error()) + h.r.JSON(w, http.StatusBadRequest, err.Error()) return } @@ -142,13 +146,13 @@ func (h *schedulerHandler) CreateScheduler(w http.ResponseWriter, r *http.Reques if errors.ErrorEqual(err, errs.ErrOptionNotExist) { collector(defaultTimeout) } else { - h.r.JSON(w, http.StatusInternalServerError, err.Error()) + h.r.JSON(w, http.StatusBadRequest, err.Error()) return } } if err := apiutil.CollectStringOption("alias", input, collector); err != nil { - h.r.JSON(w, http.StatusInternalServerError, err.Error()) + h.r.JSON(w, http.StatusBadRequest, err.Error()) return } @@ -298,7 +302,11 @@ func (h *schedulerHandler) redirectSchedulerDelete(w http.ResponseWriter, name, } resp, err := apiutil.DoDelete(h.svr.GetHTTPClient(), deleteURL) if err != nil { - h.r.JSON(w, resp.StatusCode, err.Error()) + status := http.StatusInternalServerError + if resp != nil { + status = resp.StatusCode + } + h.r.JSON(w, status, err.Error()) return } defer resp.Body.Close() @@ -373,5 +381,9 @@ func (h *schedulerConfigHandler) handleSchedulerConfig(w http.ResponseWriter, r sh.ServeHTTP(w, r) return } - h.rd.JSON(w, http.StatusNotAcceptable, err.Error()) + if err != nil { + h.rd.JSON(w, http.StatusNotAcceptable, err.Error()) + return + } + h.rd.JSON(w, http.StatusNotAcceptable, "scheduler config handler is unavailable") } From c4f92349f2868b32bbb61f49b12eabaa9b36bf9a Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Mon, 9 Mar 2026 15:14:07 +0800 Subject: [PATCH 2/2] remove unuseless change Signed-off-by: tongjian <1045931706@qq.com> --- pkg/schedule/schedulers/balance_range.go | 7 ++- pkg/schedule/schedulers/balance_range_test.go | 10 +--- pkg/schedule/schedulers/evict_leader.go | 2 +- pkg/schedule/schedulers/grant_leader.go | 28 +-------- pkg/schedule/schedulers/grant_leader_test.go | 60 ------------------- 5 files changed, 8 insertions(+), 99 deletions(-) diff --git a/pkg/schedule/schedulers/balance_range.go b/pkg/schedule/schedulers/balance_range.go index ff810ca9db6..50d74e1b19e 100644 --- a/pkg/schedule/schedulers/balance_range.go +++ b/pkg/schedule/schedulers/balance_range.go @@ -130,6 +130,10 @@ func (handler *balanceRangeSchedulerHandler) addJob(w http.ResponseWriter, r *ht handler.rd.JSON(w, http.StatusBadRequest, fmt.Sprintf("timeout:%s is invalid", timeoutStr)) return } + if timeout <= 0 { + handler.rd.JSON(w, http.StatusBadRequest, "timeout must be positive") + return + } job.Timeout = timeout } } @@ -348,9 +352,6 @@ func (job *balanceRangeSchedulerJob) expired(dur time.Duration) bool { } func (job *balanceRangeSchedulerJob) shouldFinished() bool { - if job == nil || job.Start == nil { - return true - } return time.Since(*job.Start) > job.Timeout } diff --git a/pkg/schedule/schedulers/balance_range_test.go b/pkg/schedule/schedulers/balance_range_test.go index bea3b6eaf1e..1b4d88d5d95 100644 --- a/pkg/schedule/schedulers/balance_range_test.go +++ b/pkg/schedule/schedulers/balance_range_test.go @@ -25,9 +25,9 @@ import ( "time" "github.com/stretchr/testify/require" + "github.com/unrolled/render" "github.com/pingcap/failpoint" - "github.com/unrolled/render" "github.com/tikv/pd/pkg/core" "github.com/tikv/pd/pkg/schedule/operator" @@ -549,11 +549,3 @@ func TestAddBalanceRangeJobWithInvalidFieldType(t *testing.T) { re.Equal(http.StatusBadRequest, resp.Code) re.Empty(conf.jobs) } - -func TestJobShouldFinishedWhenStartTimeIsMissing(t *testing.T) { - re := require.New(t) - job := &balanceRangeSchedulerJob{Timeout: time.Minute} - re.True(job.shouldFinished()) - job = nil - re.True(job.shouldFinished()) -} diff --git a/pkg/schedule/schedulers/evict_leader.go b/pkg/schedule/schedulers/evict_leader.go index 5adcbee79e9..2d218b00ed4 100644 --- a/pkg/schedule/schedulers/evict_leader.go +++ b/pkg/schedule/schedulers/evict_leader.go @@ -179,7 +179,7 @@ func (conf *evictLeaderSchedulerConfig) pauseLeaderTransferIfStoreNotExist(id ui if err := conf.cluster.PauseLeaderTransfer(id, constant.In); err != nil { return false, err } - return false, nil + return true, nil } func (conf *evictLeaderSchedulerConfig) resumeLeaderTransferIfExist(id uint64) { diff --git a/pkg/schedule/schedulers/grant_leader.go b/pkg/schedule/schedulers/grant_leader.go index 97c02285c2c..abd18304db9 100644 --- a/pkg/schedule/schedulers/grant_leader.go +++ b/pkg/schedule/schedulers/grant_leader.go @@ -265,32 +265,8 @@ func (handler *grantLeaderHandler) updateConfig(w http.ResponseWriter, r *http.R args = append(args, strconv.FormatUint(id, 10)) } - rangesVal, hasRanges := input["ranges"] - if hasRanges { - var ranges []string - switch val := rangesVal.(type) { - case []string: - ranges = val - case []any: - ranges = make([]string, 0, len(val)) - for _, item := range val { - s, ok := item.(string) - if !ok { - handler.config.Lock() - handler.config.cluster.ResumeLeaderTransfer(id, constant.Out) - handler.config.Unlock() - handler.rd.JSON(w, http.StatusBadRequest, errs.ErrSchedulerConfig.FastGenByArgs("ranges")) - return - } - ranges = append(ranges, s) - } - default: - handler.config.Lock() - handler.config.cluster.ResumeLeaderTransfer(id, constant.Out) - handler.config.Unlock() - handler.rd.JSON(w, http.StatusBadRequest, errs.ErrSchedulerConfig.FastGenByArgs("ranges")) - return - } + ranges, ok := (input["ranges"]).([]string) + if ok { args = append(args, ranges...) } else if exists { args = append(args, handler.config.getRanges(id)...) diff --git a/pkg/schedule/schedulers/grant_leader_test.go b/pkg/schedule/schedulers/grant_leader_test.go index 34718c4bedd..003966db73d 100644 --- a/pkg/schedule/schedulers/grant_leader_test.go +++ b/pkg/schedule/schedulers/grant_leader_test.go @@ -15,19 +15,13 @@ package schedulers import ( - "bytes" - "encoding/json" - "net/http" - "net/http/httptest" "testing" "github.com/stretchr/testify/require" - "github.com/unrolled/render" "github.com/tikv/pd/pkg/schedule/operator" "github.com/tikv/pd/pkg/schedule/types" "github.com/tikv/pd/pkg/storage" - "github.com/tikv/pd/pkg/utils/keyutil" "github.com/tikv/pd/pkg/utils/operatorutil" ) @@ -66,57 +60,3 @@ func TestGrantLeaderScheduler(t *testing.T) { ops, _ = bls.Schedule(tc, false) re.Empty(ops) } - -func TestGrantLeaderUpdateConfigWithStringArrayRanges(t *testing.T) { - re := require.New(t) - cancel, _, tc, _ := prepareSchedulersTest() - defer cancel() - - tc.AddLeaderStore(1, 0) - conf := &grantLeaderSchedulerConfig{ - schedulerConfig: &baseSchedulerConfig{}, - StoreIDWithRanges: make(map[uint64][]keyutil.KeyRange), - cluster: tc.GetBasicCluster(), - } - conf.init("grant-leader-test", storage.NewStorageWithMemoryBackend(), conf) - handler := &grantLeaderHandler{config: conf, rd: render.New(render.Options{IndentJSON: true})} - body, err := json.Marshal(map[string]any{ - "store_id": 1, - "ranges": []string{"100", "200"}, - }) - re.NoError(err) - req := httptest.NewRequest(http.MethodPost, "/config", bytes.NewReader(body)) - resp := httptest.NewRecorder() - re.NotPanics(func() { - handler.updateConfig(resp, req) - }) - re.Equal(http.StatusOK, resp.Code) - re.Len(conf.StoreIDWithRanges[1], 1) - re.Equal(keyutil.NewKeyRange("100", "200"), conf.StoreIDWithRanges[1][0]) -} - -func TestGrantLeaderUpdateConfigWithInvalidRangesType(t *testing.T) { - re := require.New(t) - cancel, _, tc, _ := prepareSchedulersTest() - defer cancel() - - tc.AddLeaderStore(1, 0) - conf := &grantLeaderSchedulerConfig{ - schedulerConfig: &baseSchedulerConfig{}, - StoreIDWithRanges: make(map[uint64][]keyutil.KeyRange), - cluster: tc.GetBasicCluster(), - } - conf.init("grant-leader-test", storage.NewStorageWithMemoryBackend(), conf) - handler := &grantLeaderHandler{config: conf, rd: render.New(render.Options{IndentJSON: true})} - body, err := json.Marshal(map[string]any{ - "store_id": 1, - "ranges": []any{"100", 200}, - }) - re.NoError(err) - req := httptest.NewRequest(http.MethodPost, "/config", bytes.NewReader(body)) - resp := httptest.NewRecorder() - re.NotPanics(func() { - handler.updateConfig(resp, req) - }) - re.Equal(http.StatusBadRequest, resp.Code) -}