pass ut

Signed-off-by: 童剑 <[email protected]>
tikv · Jan 14, 2025 · cb8a4b9 · cb8a4b9
1 parent d1da5b5
commit cb8a4b9
Show file tree

Hide file tree

Showing 8 changed files with 38 additions and 14 deletions.
diff --git a/pkg/mcs/scheduling/server/cluster.go b/pkg/mcs/scheduling/server/cluster.go
@@ -316,7 +316,8 @@ func (c *Cluster) updateScheduler() {
 		for _, scheduler := range latestSchedulersConfig {
 			schedulerType, ok := types.ConvertOldStrToType[scheduler.Type]
 			if !ok {
-				log.Warn("scheduler not found  ", zap.String("type", scheduler.Type))
+				log.Error("scheduler not found  ", zap.String("type", scheduler.Type))
+				continue
 			}
 			s, err := schedulers.CreateScheduler(
 				schedulerType,

diff --git a/pkg/schedule/operator/kind.go b/pkg/schedule/operator/kind.go
@@ -62,6 +62,7 @@ var flagToName = map[OpKind]string{
 	OpHotRegion:     "hot-region",
 	OpReplica:       "replica",
 	OpMerge:         "merge",
+	OpKeyRange:      "key-range",
 	OpRange:         "range",
 	OpWitness:       "witness",
 	OpWitnessLeader: "witness-leader",

diff --git a/pkg/schedule/schedulers/balance_key_range.go b/pkg/schedule/schedulers/balance_key_range.go
@@ -17,11 +17,7 @@ import (
 	"github.com/tikv/pd/pkg/schedule/operator"
 	"github.com/tikv/pd/pkg/schedule/plan"
 	"github.com/tikv/pd/pkg/schedule/types"
-)
-
-const (
-	// DefaultTimeout is the default balance key range scheduler timeout.
-	DefaultTimeout = 1 * time.Hour
+	"github.com/tikv/pd/pkg/utils/syncutil"
 )
 
 type balanceKeyRangeSchedulerHandler struct {
@@ -52,7 +48,8 @@ func (handler *balanceKeyRangeSchedulerHandler) listConfig(w http.ResponseWriter
 }
 
 type balanceKeyRangeSchedulerConfig struct {
-	baseDefaultSchedulerConfig
+	syncutil.RWMutex
+	schedulerConfig
 	balanceKeyRangeSchedulerParam
 }
 

diff --git a/pkg/schedule/schedulers/init.go b/pkg/schedule/schedulers/init.go
@@ -19,6 +19,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/errs"
@@ -548,7 +549,7 @@ func schedulersRegister() {
 	})
 
 	// balance key range scheduler
-	// args: [role, engine, range1, range2, ...]
+	// args: [role, engine, timeout, range1, range2, ...]
 	RegisterSliceDecoderBuilder(types.BalanceKeyRangeScheduler, func(args []string) ConfigDecoder {
 		return func(v any) error {
 			conf, ok := v.(*balanceKeyRangeSchedulerConfig)
@@ -566,22 +567,30 @@ func schedulersRegister() {
 			if err != nil {
 				return errs.ErrQueryUnescape.Wrap(err)
 			}
-			ranges, err := getKeyRanges(args[2:])
+			timeout, err := url.QueryUnescape(args[2])
+			if err != nil {
+				return errs.ErrQueryUnescape.Wrap(err)
+			}
+			duration, err := time.ParseDuration(timeout)
+			if err != nil {
+				return errs.ErrURLParse.Wrap(err)
+			}
+			ranges, err := getKeyRanges(args[3:])
 			if err != nil {
 				return err
 			}
 			conf.Ranges = ranges
 			conf.Engine = engine
 			conf.Role = role
-			conf.Timeout = DefaultTimeout
+			conf.Timeout = duration
 			return nil
 		}
 	})
 
 	RegisterScheduler(types.BalanceKeyRangeScheduler, func(opController *operator.Controller,
 		storage endpoint.ConfigStorage, decoder ConfigDecoder, _ ...func(string) error) (Scheduler, error) {
 		conf := &balanceKeyRangeSchedulerConfig{
-			baseDefaultSchedulerConfig: newBaseDefaultSchedulerConfig(),
+			schedulerConfig: newBaseDefaultSchedulerConfig(),
 		}
 		if err := decoder(conf); err != nil {
 			return nil, err

diff --git a/pkg/schedule/schedulers/scheduler.go b/pkg/schedule/schedulers/scheduler.go
@@ -157,9 +157,7 @@ func CreateScheduler(
 	removeSchedulerCb ...func(string) error,
 ) (Scheduler, error) {
 	fn, ok := schedulerMap[typ]
-	log.Info("create scheduler", zap.Any("typ", typ))
 	if !ok {
-		log.Warn("create scheduler not found", zap.Any("typ", typ))
 		return nil, errs.ErrSchedulerCreateFuncNotRegistered.FastGenByArgs(typ)
 	}
 

diff --git a/server/api/scheduler.go b/server/api/scheduler.go
@@ -113,6 +113,16 @@ func (h *schedulerHandler) CreateScheduler(w http.ResponseWriter, r *http.Reques
 			h.r.JSON(w, http.StatusInternalServerError, err.Error())
 			return
 		}
+		defaultTimeout := "1h"
+		if err := apiutil.CollectStringOption("timeout", input, collector); err != nil {
+			if errors.ErrorEqual(err, errs.ErrOptionNotExist) {
+				collector(defaultTimeout)
+			} else {
+				h.r.JSON(w, http.StatusInternalServerError, err.Error())
+				return
+			}
+		}
+
 		if err := apiutil.CollectEscapeStringOption("start_key", input, collector); err != nil {
 			h.r.JSON(w, http.StatusInternalServerError, err.Error())
 			return

diff --git a/server/cluster/cluster_test.go b/server/cluster/cluster_test.go
@@ -3212,9 +3212,16 @@ func TestAddScheduler(t *testing.T) {
 	_, err = schedulers.CreateScheduler(types.BalanceKeyRangeScheduler, oc, storage.NewStorageWithMemoryBackend(), schedulers.ConfigSliceDecoder(types.BalanceKeyRangeScheduler, []string{}), controller.RemoveScheduler)
 	re.Error(err)
 
-	gls, err = schedulers.CreateScheduler(types.BalanceKeyRangeScheduler, oc, storage.NewStorageWithMemoryBackend(), schedulers.ConfigSliceDecoder(types.BalanceKeyRangeScheduler, []string{"leaner", "tiflash", "100", "200"}), controller.RemoveScheduler)
+	gls, err = schedulers.CreateScheduler(types.BalanceKeyRangeScheduler, oc, storage.NewStorageWithMemoryBackend(), schedulers.ConfigSliceDecoder(types.BalanceKeyRangeScheduler, []string{"learner", "tiflash", "1h", "100", "200"}), controller.RemoveScheduler)
 	re.NoError(err)
 	re.NoError(controller.AddScheduler(gls))
+	conf, err = gls.EncodeConfig()
+	re.NoError(err)
+	data = make(map[string]any)
+	re.NoError(json.Unmarshal(conf, &data))
+	re.Equal("learner", data["role"])
+	re.Equal("tiflash", data["engine"])
+	re.Equal(float64(time.Hour.Nanoseconds()), data["timeout"])
 
 	hb, err := schedulers.CreateScheduler(types.BalanceHotRegionScheduler, oc, storage.NewStorageWithMemoryBackend(), schedulers.ConfigJSONDecoder([]byte("{}")))
 	re.NoError(err)

diff --git a/tools/pd-ctl/tests/scheduler/scheduler_test.go b/tools/pd-ctl/tests/scheduler/scheduler_test.go
@@ -551,6 +551,7 @@ func (suite *schedulerTestSuite) checkSchedulerConfig(cluster *pdTests.TestClust
 	mustExec(re, cmd, []string{"-u", pdAddr, "scheduler", "config", "balance-key-range-scheduler", "show"}, &conf)
 	re.Equal("learner", conf["role"])
 	re.Equal("tiflash", conf["engine"])
+	re.Equal(float64(time.Hour.Nanoseconds()), conf["timeout"])
 	ranges := conf["ranges"].([]any)[0].(map[string]any)
 	re.Equal(base64.StdEncoding.EncodeToString([]byte("a")), ranges["start-key"])
 	re.Equal(base64.StdEncoding.EncodeToString([]byte("b")), ranges["end-key"])