diff --git a/Makefile b/Makefile index 205896c377a..f24dcc62f15 100644 --- a/Makefile +++ b/Makefile @@ -80,7 +80,7 @@ BUILD_BIN_PATH := $(ROOT_PATH)/bin build: pd-server pd-ctl pd-recover -tools: pd-tso-bench pd-heartbeat-bench regions-dump stores-dump pd-api-bench pd-ut +tools: pd-dev PD_SERVER_DEP := ifeq ($(SWAGGER), 1) @@ -110,28 +110,14 @@ pd-server-basic: # Tools pd-ctl: cd tools && GOEXPERIMENT=$(BUILD_GOEXPERIMENT) CGO_ENABLED=$(BUILD_TOOL_CGO_ENABLED) go build -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/pd-ctl pd-ctl/main.go -pd-tso-bench: - cd tools && CGO_ENABLED=0 go build -o $(BUILD_BIN_PATH)/pd-tso-bench pd-tso-bench/main.go -pd-api-bench: - cd tools && CGO_ENABLED=0 go build -o $(BUILD_BIN_PATH)/pd-api-bench pd-api-bench/main.go +pd-dev: pd-xprog + cd tools && GOEXPERIMENT=$(BUILD_GOEXPERIMENT) CGO_ENABLED=$(BUILD_TOOL_CGO_ENABLED) go build -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/pd-dev pd-dev/main.go pd-recover: cd tools && GOEXPERIMENT=$(BUILD_GOEXPERIMENT) CGO_ENABLED=$(BUILD_TOOL_CGO_ENABLED) go build -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/pd-recover pd-recover/main.go -pd-analysis: - cd tools && CGO_ENABLED=0 go build -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/pd-analysis pd-analysis/main.go -pd-heartbeat-bench: - cd tools && CGO_ENABLED=0 go build -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/pd-heartbeat-bench pd-heartbeat-bench/main.go -simulator: - cd tools && CGO_ENABLED=0 go build -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/pd-simulator pd-simulator/main.go -regions-dump: - cd tools && CGO_ENABLED=0 go build -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/regions-dump regions-dump/main.go -stores-dump: - cd tools && CGO_ENABLED=0 go build -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/stores-dump stores-dump/main.go -pd-ut: pd-xprog - cd tools && GOEXPERIMENT=$(BUILD_GOEXPERIMENT) CGO_ENABLED=$(BUILD_TOOL_CGO_ENABLED) go build -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/pd-ut pd-ut/ut.go pd-xprog: - cd tools && GOEXPERIMENT=$(BUILD_GOEXPERIMENT) CGO_ENABLED=$(BUILD_TOOL_CGO_ENABLED) go build -tags xprog -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/xprog pd-ut/xprog.go + cd tools && GOEXPERIMENT=$(BUILD_GOEXPERIMENT) CGO_ENABLED=$(BUILD_TOOL_CGO_ENABLED) go build -tags xprog -gcflags '$(GCFLAGS)' -ldflags '$(LDFLAGS)' -o $(BUILD_BIN_PATH)/xprog pd-dev/pd-ut/xprog.go -.PHONY: pd-ctl pd-tso-bench pd-recover pd-analysis pd-heartbeat-bench simulator regions-dump stores-dump pd-api-bench pd-ut +.PHONY: pd-ctl pd-recover pd-dev #### Docker image #### diff --git a/tools/pd-analysis/main.go b/tools/pd-analysis/main.go deleted file mode 100644 index 510448ddddd..00000000000 --- a/tools/pd-analysis/main.go +++ /dev/null @@ -1,82 +0,0 @@ -// Copyright 2019 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 main - -import ( - "flag" - "os" - - "github.com/pingcap/log" - "github.com/tikv/pd/tools/pd-analysis/analysis" - "go.uber.org/zap" -) - -var ( - input = flag.String("input", "", "input pd log file, required") - output = flag.String("output", "", "output file, default output to stdout") - logLevel = flag.String("logLevel", "info", "log level, default info") - style = flag.String("style", "", "analysis style, e.g. transfer-counter") - operator = flag.String("operator", "", "operator style, e.g. balance-region, balance-leader, transfer-hot-read-leader, move-hot-read-region, transfer-hot-write-leader, move-hot-write-region") - start = flag.String("start", "", "start time, e.g. 2019/09/10 12:20:07, default: total file") - end = flag.String("end", "", "end time, e.g. 2019/09/10 14:20:07, default: total file") -) - -// Logger is the global logger used for simulator. -var Logger *zap.Logger - -// InitLogger initializes the Logger with -log level. -func InitLogger(l string) { - conf := &log.Config{Level: l, File: log.FileLogConfig{}} - lg, _, _ := log.InitLogger(conf) - Logger = lg -} - -func main() { - flag.Parse() - InitLogger(*logLevel) - analysis.GetTransferCounter().Init(0, 0) - if *input == "" { - Logger.Fatal("need to specify one input pd log") - } - if *output != "" { - f, err := os.OpenFile(*output, os.O_WRONLY|os.O_CREATE|os.O_SYNC|os.O_APPEND, 0600) - if err != nil { - Logger.Fatal(err.Error()) - } else { - os.Stdout = f - } - } - - switch *style { - case "transfer-counter": - { - if *operator == "" { - Logger.Fatal("need to specify one operator") - } - r, err := analysis.GetTransferCounter().CompileRegex(*operator) - if err != nil { - Logger.Fatal(err.Error()) - } - err = analysis.GetTransferCounter().ParseLog(*input, *start, *end, analysis.DefaultLayout, r) - if err != nil { - Logger.Fatal(err.Error()) - } - analysis.GetTransferCounter().PrintResult() - break - } - default: - Logger.Fatal("style is not exist") - } -} diff --git a/tools/pd-dev/README.md b/tools/pd-dev/README.md new file mode 100644 index 00000000000..c372e6bfdcd --- /dev/null +++ b/tools/pd-dev/README.md @@ -0,0 +1,70 @@ +# pd-dev + +pd-dev is a tool for developing and testing PD. It provides the following functions: + +- [pd-api-bench](./pd-api-bench/README.md): A tool for testing the performance of PD's API. +- [pd-tso-bench](./pd-tso-bench/README.md): A tool for testing the performance of PD's TSO. +- [pd-heartbeat-bench](./pd-heartbeat-bench/README.md): A tool for testing the performance of PD's heartbeat. +- [pd-simulator](./pd-simulator/README.md): A tool for simulating the PD cluster. +- [regions-dump](./regions-dump/README.md): A tool for dumping the region information of the PD cluster. +- [stores-dump](./stores-dump/README.md): A tool for dumping the store information of the PD cluster. + +## Build + +1. [Go](https://golang.org/) Version 1.21 or later +2. In the root directory of the [PD project](https://github.com/tikv/pd), use the `make pd-dev` command to compile and generate `bin/pd-dev` + +## Usage + +This section describes how to use the `pd-dev` tool. + +### Cases + +Please read related README files for more details, we support the following cases: + +`./pd-dev --mode api` + +`./pd-dev --mode tso` + +`./pd-dev --mode heartbeat` + +`./pd-dev --mode simulator` + +`./pd-dev --mode regions-dump` + +`./pd-dev --mode stores-dump` + +`./pd-dev --mode analysis` + +`./pd-dev --mode backup` + +`./pd-dev --mode ut` + +### flag description + +--pd string +> pd address (default "127.0.0.1:2379") + +--cacert string +> path of file that contains list of trusted SSL CAs + +--cert string +> path of file that contains X509 certificate in PEM format + +--key string +> path of file that contains X509 key in PEM format + +--config string +> path of configuration file + +### TLS + +You can use the following command to generate a certificate for testing TLS: + +```shell +mkdir cert +./cert_opt.sh generate cert +./bin/pd-dev --mode api -http-cases GetRegionStatus-1+1,GetMinResolvedTS-1+1 -client 1 -debug true -cacert ./cert/ca.pem -cert ./cert/pd-server.pem -key ./cert/pd-server-key.pem +./cert_opt.sh cleanup cert +rm -rf cert +``` diff --git a/tools/pd-dev/main.go b/tools/pd-dev/main.go new file mode 100644 index 00000000000..63dfb2fa83d --- /dev/null +++ b/tools/pd-dev/main.go @@ -0,0 +1,100 @@ +// Copyright 2024 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 main + +import ( + "context" + "os" + "os/signal" + "syscall" + + "github.com/pingcap/log" + flag "github.com/spf13/pflag" + analysis "github.com/tikv/pd/tools/pd-dev/pd-analysis" + api "github.com/tikv/pd/tools/pd-dev/pd-api-bench" + backup "github.com/tikv/pd/tools/pd-dev/pd-backup" + heartbeat "github.com/tikv/pd/tools/pd-dev/pd-heartbeat-bench" + simulator "github.com/tikv/pd/tools/pd-dev/pd-simulator" + tso "github.com/tikv/pd/tools/pd-dev/pd-tso-bench" + ut "github.com/tikv/pd/tools/pd-dev/pd-ut" + region "github.com/tikv/pd/tools/pd-dev/regions-dump" + store "github.com/tikv/pd/tools/pd-dev/stores-dump" + "go.uber.org/zap" +) + +func switchDevelopmentMode(ctx context.Context, mode string) { + log.Info("pd-dev start", zap.String("mode", mode)) + switch mode { + case "api": + api.Run(ctx) + case "heartbeat": + heartbeat.Run(ctx) + case "tso": + tso.Run(ctx) + case "simulator": + simulator.Run(ctx) + case "regions-dump": + region.Run() + case "stores-dump": + store.Run() + case "analysis": + analysis.Run() + case "backup": + backup.Run() + case "ut": + ut.Run() + default: + log.Fatal("please input a mode to run, or provide a config file to run all modes") + } +} + +func main() { + ctx, cancel := context.WithCancel(context.Background()) + + sc := make(chan os.Signal, 1) + signal.Notify(sc, + syscall.SIGHUP, + syscall.SIGINT, + syscall.SIGTERM, + syscall.SIGQUIT) + var sig os.Signal + go func() { + sig = <-sc + cancel() + }() + + // parse first argument + if len(os.Args) < 2 { + log.Fatal("please input a mode to run, or provide a config file to run all modes") + } + + var mode string + fs := flag.NewFlagSet("pd-dev", flag.ContinueOnError) + fs.StringVar(&mode, "mode", "", "mode to run") + fs.Parse(os.Args[1:]) + switchDevelopmentMode(ctx, mode) + + log.Info("pd-dev Exit") + switch sig { + case syscall.SIGTERM: + exit(0) + default: + exit(1) + } +} + +func exit(code int) { + os.Exit(code) +} diff --git a/tools/pd-analysis/analysis/parse_log.go b/tools/pd-dev/pd-analysis/analysis/parse_log.go similarity index 100% rename from tools/pd-analysis/analysis/parse_log.go rename to tools/pd-dev/pd-analysis/analysis/parse_log.go diff --git a/tools/pd-analysis/analysis/parse_log_test.go b/tools/pd-dev/pd-analysis/analysis/parse_log_test.go similarity index 100% rename from tools/pd-analysis/analysis/parse_log_test.go rename to tools/pd-dev/pd-analysis/analysis/parse_log_test.go diff --git a/tools/pd-analysis/analysis/transfer_counter.go b/tools/pd-dev/pd-analysis/analysis/transfer_counter.go similarity index 98% rename from tools/pd-analysis/analysis/transfer_counter.go rename to tools/pd-dev/pd-analysis/analysis/transfer_counter.go index 8f472ae1e03..e337d5a6077 100644 --- a/tools/pd-analysis/analysis/transfer_counter.go +++ b/tools/pd-dev/pd-analysis/analysis/transfer_counter.go @@ -74,7 +74,7 @@ func (c *TransferCounter) Init(storeNum, regionNum int) { c.loopResultCount = c.loopResultCount[:0] } -// AddTarget is be used to add target of edge in graph mat. +// AddTarget be used to add target of edge in graph mat. // Firstly add a new peer and then delete the old peer of the scheduling, // So in the statistics, also firstly add the target and then add the source. func (c *TransferCounter) AddTarget(regionID, targetStoreID uint64) { @@ -83,7 +83,7 @@ func (c *TransferCounter) AddTarget(regionID, targetStoreID uint64) { c.regionMap[regionID] = targetStoreID } -// AddSource is be used to add source of edge in graph mat. +// AddSource be used to add source of edge in graph mat. func (c *TransferCounter) AddSource(regionID, sourceStoreID uint64) { c.mutex.Lock() defer c.mutex.Unlock() diff --git a/tools/pd-analysis/analysis/transfer_counter_test.go b/tools/pd-dev/pd-analysis/analysis/transfer_counter_test.go similarity index 100% rename from tools/pd-analysis/analysis/transfer_counter_test.go rename to tools/pd-dev/pd-analysis/analysis/transfer_counter_test.go diff --git a/tools/pd-dev/pd-analysis/config.go b/tools/pd-dev/pd-analysis/config.go new file mode 100644 index 00000000000..1a34786cb1d --- /dev/null +++ b/tools/pd-dev/pd-analysis/config.go @@ -0,0 +1,51 @@ +// Copyright 2024 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 analysis + +import ( + flag "github.com/spf13/pflag" + "github.com/tikv/pd/tools/pd-dev/util" +) + +// Config is the heartbeat-bench configuration. +type Config struct { + *util.GeneralConfig + Input string `toml:"input" json:"input"` + Output string `toml:"output" json:"output"` + Style string `toml:"style" json:"style"` + Operator string `toml:"operator" json:"operator"` + Start string `toml:"start" json:"start"` + End string `toml:"end" json:"end"` +} + +// NewConfig return a set of settings. +func newConfig() *Config { + cfg := &Config{ + GeneralConfig: &util.GeneralConfig{}, + } + cfg.FlagSet = flag.NewFlagSet("pd-analysis", flag.ContinueOnError) + fs := cfg.FlagSet + cfg.GeneralConfig = util.NewGeneralConfig(fs) + fs.ParseErrorsWhitelist.UnknownFlags = true + + fs.StringVar(&cfg.Input, "input", "", "input pd log file, required") + fs.StringVar(&cfg.Output, "output", "", "output file, default output to stdout") + fs.StringVar(&cfg.Style, "style", "", "analysis style, e.g. transfer-counter") + fs.StringVar(&cfg.Operator, "operator", "", "operator style, e.g. balance-region, balance-leader, transfer-hot-read-leader, move-hot-read-region, transfer-hot-write-leader, move-hot-write-region") + fs.StringVar(&cfg.Start, "start", "", "start time, e.g. 2019/09/10 12:20:07, default: total file") + fs.StringVar(&cfg.End, "end", "", "end time, e.g. 2019/09/10 14:20:07, default: total file") + + return cfg +} diff --git a/tools/pd-dev/pd-analysis/main.go b/tools/pd-dev/pd-analysis/main.go new file mode 100644 index 00000000000..b0430af7f19 --- /dev/null +++ b/tools/pd-dev/pd-analysis/main.go @@ -0,0 +1,70 @@ +// Copyright 2019 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 analysis + +import ( + "flag" + "os" + + "github.com/pingcap/errors" + "github.com/pingcap/log" + "github.com/tikv/pd/tools/pd-dev/pd-analysis/analysis" + "go.uber.org/zap" +) + +func Run() { + cfg := newConfig() + err := cfg.Parse(os.Args[1:]) + switch errors.Cause(err) { + case nil: + case flag.ErrHelp: + os.Exit(0) + default: + log.Fatal("parse cmd flags error", zap.Error(err)) + } + + if cfg.Input == "" { + log.Fatal("need to specify one input pd log") + } + analysis.GetTransferCounter().Init(0, 0) + if cfg.Output != "" { + f, err := os.OpenFile(cfg.Output, os.O_WRONLY|os.O_CREATE|os.O_SYNC|os.O_APPEND, 0600) + if err != nil { + log.Fatal(err.Error()) + } + os.Stdout = f + } + + switch cfg.Style { + case "transfer-counter": + { + if cfg.Operator == "" { + log.Fatal("need to specify one operator") + } + r, err := analysis.GetTransferCounter().CompileRegex(cfg.Operator) + if err != nil { + log.Fatal(err.Error()) + } + err = analysis.GetTransferCounter().ParseLog(cfg.Input, cfg.Start, cfg.End, analysis.DefaultLayout, r) + if err != nil { + log.Fatal(err.Error()) + } + analysis.GetTransferCounter().PrintResult() + break + } + default: + log.Fatal("style is not exist") + } +} diff --git a/tools/pd-api-bench/README.md b/tools/pd-dev/pd-api-bench/README.md similarity index 64% rename from tools/pd-api-bench/README.md rename to tools/pd-dev/pd-api-bench/README.md index 019dbd5e6a3..085d31b4439 100644 --- a/tools/pd-api-bench/README.md +++ b/tools/pd-dev/pd-api-bench/README.md @@ -29,36 +29,24 @@ The api bench cases we support are as follows: ### Flags description --cacert string -> path of file that contains list of trusted SSL CAs - --cert string -> path of file that contains X509 certificate in PEM format - --key string -> path of file that contains X509 key in PEM format - --client int +--client int > the client number (default 1) --pd string -> pd address (default "127.0.0.1:2379") - --http-cases +--http-cases > HTTP api bench cases list. Multiple cases are separated by commas. Such as `-http-cases GetRegionStatus,GetMinResolvedTS` > You can set qps with '-' and set burst with '+' such as `-http-cases GetRegionStatus-100+1` --grpc-cases +--grpc-cases > gRPC api bench cases list. Multiple cases are separated by commas. Such as `-grpc-cases GetRegion` > You can set qps with '-' and set burst with '+' such as `-grpc-cases GetRegion-1000000+10` --qps +--qps > the qps of request (default 1000). It will set qps for all cases except those in which qps is specified separately --burst +--burst > the burst of request (default 1). It will set burst for all cases except those in which burst is specified separately --debug +--debug > print the output of api response for debug ### Run Shell @@ -66,17 +54,5 @@ The api bench cases we support are as follows: You can run shell as follows. ```shell -go run main.go -http-cases GetRegionStatus-1+1,GetMinResolvedTS-1+1 -client 1 -debug true -``` - -### TLS - -You can use the following command to generate a certificate for testing TLS: - -```shell -mkdir cert -./cert_opt.sh generate cert -go run main.go -http-cases GetRegionStatus-1+1,GetMinResolvedTS-1+1 -client 1 -debug true -cacert ./cert/ca.pem -cert ./cert/pd-server.pem -key ./cert/pd-server-key.pem -./cert_opt.sh cleanup cert -rm -rf cert +./bin/pd-dev --mode api --http-cases GetRegionStatus-1+1,GetMinResolvedTS-1+1 --client 1 --debug true ``` diff --git a/tools/pd-api-bench/cases/cases.go b/tools/pd-dev/pd-api-bench/cases/cases.go similarity index 100% rename from tools/pd-api-bench/cases/cases.go rename to tools/pd-dev/pd-api-bench/cases/cases.go diff --git a/tools/pd-api-bench/cases/controller.go b/tools/pd-dev/pd-api-bench/cases/controller.go similarity index 100% rename from tools/pd-api-bench/cases/controller.go rename to tools/pd-dev/pd-api-bench/cases/controller.go diff --git a/tools/pd-api-bench/cert_opt.sh b/tools/pd-dev/pd-api-bench/cert_opt.sh similarity index 100% rename from tools/pd-api-bench/cert_opt.sh rename to tools/pd-dev/pd-api-bench/cert_opt.sh diff --git a/tools/pd-dev/pd-api-bench/config/config.go b/tools/pd-dev/pd-api-bench/config/config.go new file mode 100644 index 00000000000..9ae52db3686 --- /dev/null +++ b/tools/pd-dev/pd-api-bench/config/config.go @@ -0,0 +1,78 @@ +// Copyright 2024 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 ( + "github.com/pingcap/log" + flag "github.com/spf13/pflag" + "github.com/tikv/pd/tools/pd-dev/pd-api-bench/cases" + "github.com/tikv/pd/tools/pd-dev/util" + "go.uber.org/zap" +) + +// Config is the heartbeat-bench configuration. +type Config struct { + *util.GeneralConfig + Client int64 `toml:"client" json:"client"` + + QPS int64 + Burst int64 + HTTPCases string + GRPCCases string + + // only for init + HTTP map[string]cases.Config `toml:"http" json:"http"` + GRPC map[string]cases.Config `toml:"grpc" json:"grpc"` + ETCD map[string]cases.Config `toml:"etcd" json:"etcd"` +} + +// NewConfig return a set of settings. +func NewConfig() *Config { + cfg := &Config{ + GeneralConfig: &util.GeneralConfig{}, + } + cfg.FlagSet = flag.NewFlagSet("api-bench", flag.ContinueOnError) + fs := cfg.FlagSet + cfg.GeneralConfig = util.NewGeneralConfig(fs) + fs.ParseErrorsWhitelist.UnknownFlags = true + fs.Int64Var(&cfg.QPS, "qps", 1, "qps") + fs.Int64Var(&cfg.Burst, "burst", 1, "burst") + fs.StringVar(&cfg.HTTPCases, "http-cases", "", "http api cases") + fs.StringVar(&cfg.GRPCCases, "grpc-cases", "", "grpc cases") + fs.Int64Var(&cfg.Client, "client", 1, "client number") + return cfg +} + +// InitCoordinator set case config from config itself. +func (c *Config) InitCoordinator(co *cases.Coordinator) { + for name, cfg := range c.HTTP { + err := co.SetHTTPCase(name, &cfg) + if err != nil { + log.Error("create HTTP case failed", zap.Error(err)) + } + } + for name, cfg := range c.GRPC { + err := co.SetGRPCCase(name, &cfg) + if err != nil { + log.Error("create gRPC case failed", zap.Error(err)) + } + } + for name, cfg := range c.ETCD { + err := co.SetETCDCase(name, &cfg) + if err != nil { + log.Error("create etcd case failed", zap.Error(err)) + } + } +} diff --git a/tools/pd-api-bench/config/simconfig.toml b/tools/pd-dev/pd-api-bench/config/simconfig.toml similarity index 100% rename from tools/pd-api-bench/config/simconfig.toml rename to tools/pd-dev/pd-api-bench/config/simconfig.toml diff --git a/tools/pd-api-bench/main.go b/tools/pd-dev/pd-api-bench/main.go similarity index 71% rename from tools/pd-api-bench/main.go rename to tools/pd-dev/pd-api-bench/main.go index f9feeeea580..f7c24dc99c3 100644 --- a/tools/pd-api-bench/main.go +++ b/tools/pd-dev/pd-api-bench/main.go @@ -12,17 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package apibench import ( "context" - "crypto/tls" "net/http" "os" - "os/signal" "strconv" "strings" - "syscall" "time" "github.com/gin-contrib/cors" @@ -31,86 +28,37 @@ import ( "github.com/gin-gonic/gin" "github.com/pingcap/log" "github.com/pkg/errors" - "github.com/prometheus/client_golang/prometheus" flag "github.com/spf13/pflag" pd "github.com/tikv/pd/client" pdHttp "github.com/tikv/pd/client/http" "github.com/tikv/pd/client/tlsutil" "github.com/tikv/pd/pkg/mcs/utils" "github.com/tikv/pd/pkg/utils/logutil" - "github.com/tikv/pd/tools/pd-api-bench/cases" - "github.com/tikv/pd/tools/pd-api-bench/config" + "github.com/tikv/pd/tools/pd-dev/pd-api-bench/cases" + "github.com/tikv/pd/tools/pd-dev/pd-api-bench/config" + "github.com/tikv/pd/tools/pd-dev/pd-api-bench/metrics" + "github.com/tikv/pd/tools/pd-dev/util" "go.etcd.io/etcd/clientv3" "go.uber.org/zap" "google.golang.org/grpc" "google.golang.org/grpc/keepalive" ) -var ( - qps, burst int64 - httpCases, gRPCCases string -) - -var ( - pdAPIExecutionHistogram = prometheus.NewHistogramVec( - prometheus.HistogramOpts{ - Namespace: "pd", - Subsystem: "api_bench", - Name: "pd_api_execution_duration_seconds", - Help: "Bucketed histogram of all pd api execution time (s)", - Buckets: prometheus.ExponentialBuckets(0.001, 2, 20), // 1ms ~ 524s - }, []string{"type"}) - - pdAPIRequestCounter = prometheus.NewCounterVec( - prometheus.CounterOpts{ - Namespace: "pd", - Subsystem: "api_bench", - Name: "pd_api_request_total", - Help: "Counter of the pd http api requests", - }, []string{"type", "result"}) -) - -func main() { - prometheus.MustRegister(pdAPIExecutionHistogram) - prometheus.MustRegister(pdAPIRequestCounter) - - ctx, cancel := context.WithCancel(context.Background()) - flagSet := flag.NewFlagSet("api-bench", flag.ContinueOnError) - flagSet.ParseErrorsWhitelist.UnknownFlags = true - flagSet.Int64Var(&qps, "qps", 1, "qps") - flagSet.Int64Var(&burst, "burst", 1, "burst") - flagSet.StringVar(&httpCases, "http-cases", "", "http api cases") - flagSet.StringVar(&gRPCCases, "grpc-cases", "", "grpc cases") - cfg := config.NewConfig(flagSet) - err := cfg.Parse(os.Args[1:]) +func Run(ctx context.Context) { defer logutil.LogPanic() + metrics.RegisterMetrics() + cfg := config.NewConfig() + err := cfg.Parse(os.Args[1:]) switch errors.Cause(err) { case nil: case flag.ErrHelp: - exit(0) + os.Exit(0) default: log.Fatal("parse cmd flags error", zap.Error(err)) } - err = logutil.SetupLogger(cfg.Log, &cfg.Logger, &cfg.LogProps) - if err == nil { - log.ReplaceGlobals(cfg.Logger, cfg.LogProps) - } else { - log.Fatal("initialize logger error", zap.Error(err)) - } - sc := make(chan os.Signal, 1) - signal.Notify(sc, - syscall.SIGHUP, - syscall.SIGINT, - syscall.SIGTERM, - syscall.SIGQUIT) - - var sig os.Signal - go func() { - sig = <-sc - cancel() - }() + // set client if cfg.Client == 0 { log.Error("concurrency == 0, exit") return @@ -127,7 +75,9 @@ func main() { httpClis := make([]pdHttp.Client, cfg.Client) for i := int64(0); i < cfg.Client; i++ { sd := pdClis[i].GetServiceDiscovery() - httpClis[i] = pdHttp.NewClientWithServiceDiscovery("tools-api-bench", sd, pdHttp.WithTLSConfig(loadTLSConfig(cfg)), pdHttp.WithMetrics(pdAPIRequestCounter, pdAPIExecutionHistogram)) + httpClis[i] = pdHttp.NewClientWithServiceDiscovery("tools-api-bench", sd, + pdHttp.WithTLSConfig(util.LoadTLSConfig(cfg.GeneralConfig)), + pdHttp.WithMetrics(metrics.PDAPIRequestCounter, metrics.PDAPIExecutionHistogram)) } err = cases.InitCluster(ctx, pdClis[0], httpClis[0]) if err != nil { @@ -136,17 +86,17 @@ func main() { coordinator := cases.NewCoordinator(ctx, httpClis, pdClis, etcdClis) - hcaseStr := strings.Split(httpCases, ",") + hcaseStr := strings.Split(cfg.HTTPCases, ",") for _, str := range hcaseStr { - name, cfg := parseCaseNameAndConfig(str) + name, cfg := parseCaseNameAndConfig(cfg.QPS, cfg.Burst, str) if len(name) == 0 { continue } coordinator.SetHTTPCase(name, cfg) } - gcaseStr := strings.Split(gRPCCases, ",") + gcaseStr := strings.Split(cfg.GRPCCases, ",") for _, str := range gcaseStr { - name, cfg := parseCaseNameAndConfig(str) + name, cfg := parseCaseNameAndConfig(cfg.QPS, cfg.Burst, str) if len(name) == 0 { continue } @@ -154,9 +104,12 @@ func main() { } cfg.InitCoordinator(coordinator) - go runHTTPServer(cfg, coordinator) + srv := runHTTPServer(cfg, coordinator) <-ctx.Done() + if err := srv.Shutdown(context.Background()); err != nil { + log.Fatal("server shutdown error", zap.Error(err)) + } for _, cli := range pdClis { cli.Close() } @@ -166,20 +119,10 @@ func main() { for _, cli := range etcdClis { cli.Close() } - log.Info("Exit") - switch sig { - case syscall.SIGTERM: - exit(0) - default: - exit(1) - } + log.Info("API bench exit") } -func exit(code int) { - os.Exit(code) -} - -func parseCaseNameAndConfig(str string) (string, *cases.Config) { +func parseCaseNameAndConfig(qps int64, burst int64, str string) (string, *cases.Config) { var err error cfg := &cases.Config{} name := "" @@ -220,7 +163,7 @@ func parseCaseNameAndConfig(str string) (string, *cases.Config) { return name, cfg } -func runHTTPServer(cfg *config.Config, co *cases.Coordinator) { +func runHTTPServer(cfg *config.Config, co *cases.Coordinator) *http.Server { gin.SetMode(gin.ReleaseMode) engine := gin.New() engine.Use(gin.Recovery()) @@ -341,7 +284,15 @@ func runHTTPServer(cfg *config.Config, co *cases.Coordinator) { } c.IndentedJSON(http.StatusOK, cfg) }) - engine.Run(cfg.StatusAddr) + + srv := &http.Server{Addr: cfg.StatusAddr, Handler: engine.Handler(), ReadHeaderTimeout: 3 * time.Second} + go func() { + if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + log.Fatal("server listen error", zap.Error(err)) + } + }() + + return srv } const ( @@ -362,7 +313,7 @@ func newEtcdClient(cfg *config.Config) *clientv3.Client { return nil } clientConfig := clientv3.Config{ - Endpoints: []string{cfg.PDAddr}, + Endpoints: []string{cfg.PDAddrs}, DialTimeout: keepaliveTimeout, TLS: tlsCfg, LogConfig: &lgc, @@ -376,7 +327,7 @@ func newEtcdClient(cfg *config.Config) *clientv3.Client { // newPDClient returns a pd client. func newPDClient(ctx context.Context, cfg *config.Config) pd.Client { - addrs := []string{cfg.PDAddr} + addrs := []string{cfg.PDAddrs} pdCli, err := pd.NewClientWithContext(ctx, addrs, pd.SecurityOption{ CAPath: cfg.CaPath, CertPath: cfg.CertPath, @@ -393,32 +344,3 @@ func newPDClient(ctx context.Context, cfg *config.Config) pd.Client { } return pdCli } - -func loadTLSConfig(cfg *config.Config) *tls.Config { - if len(cfg.CaPath) == 0 { - return nil - } - caData, err := os.ReadFile(cfg.CaPath) - if err != nil { - log.Error("fail to read ca file", zap.Error(err)) - } - certData, err := os.ReadFile(cfg.CertPath) - if err != nil { - log.Error("fail to read cert file", zap.Error(err)) - } - keyData, err := os.ReadFile(cfg.KeyPath) - if err != nil { - log.Error("fail to read key file", zap.Error(err)) - } - - tlsConf, err := tlsutil.TLSConfig{ - SSLCABytes: caData, - SSLCertBytes: certData, - SSLKEYBytes: keyData, - }.ToTLSConfig() - if err != nil { - log.Fatal("failed to load tlc config", zap.Error(err)) - } - - return tlsConf -} diff --git a/tools/pd-dev/pd-api-bench/metrics/metric.go b/tools/pd-dev/pd-api-bench/metrics/metric.go new file mode 100644 index 00000000000..6c4a78787d4 --- /dev/null +++ b/tools/pd-dev/pd-api-bench/metrics/metric.go @@ -0,0 +1,41 @@ +// Copyright 2024 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 metrics + +import "github.com/prometheus/client_golang/prometheus" + +var ( + PDAPIExecutionHistogram = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: "pd", + Subsystem: "api_bench", + Name: "pd_api_execution_duration_seconds", + Help: "Bucketed histogram of all pd api execution time (s)", + Buckets: prometheus.ExponentialBuckets(0.001, 2, 20), // 1ms ~ 524s + }, []string{"type"}) + + PDAPIRequestCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: "pd", + Subsystem: "api_bench", + Name: "pd_api_request_total", + Help: "Counter of the pd http api requests", + }, []string{"type", "result"}) +) + +func RegisterMetrics() { + prometheus.MustRegister(PDAPIExecutionHistogram) + prometheus.MustRegister(PDAPIRequestCounter) +} diff --git a/tools/pd-dev/pd-backup/config.go b/tools/pd-dev/pd-backup/config.go new file mode 100644 index 00000000000..f14c8e58ac1 --- /dev/null +++ b/tools/pd-dev/pd-backup/config.go @@ -0,0 +1,42 @@ +// Copyright 2024 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 backup + +import ( + flag "github.com/spf13/pflag" + "github.com/tikv/pd/tools/pd-dev/util" +) + +// Config is the heartbeat-bench configuration. +type Config struct { + *util.GeneralConfig + + FilePath string `toml:"file" json:"file"` +} + +// NewConfig return a set of settings. +func newConfig() *Config { + cfg := &Config{ + GeneralConfig: &util.GeneralConfig{}, + } + cfg.FlagSet = flag.NewFlagSet("pd-backup", flag.ContinueOnError) + fs := cfg.FlagSet + cfg.GeneralConfig = util.NewGeneralConfig(fs) + fs.ParseErrorsWhitelist.UnknownFlags = true + + fs.StringVar(&cfg.FilePath, "file", "backup.json", "backup file path and name") + + return cfg +} diff --git a/tools/pd-backup/main.go b/tools/pd-dev/pd-backup/main.go similarity index 61% rename from tools/pd-backup/main.go rename to tools/pd-dev/pd-backup/main.go index c424113e451..9600e258daf 100644 --- a/tools/pd-backup/main.go +++ b/tools/pd-dev/pd-backup/main.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package backup import ( "flag" @@ -21,26 +21,30 @@ import ( "strings" "time" - "github.com/tikv/pd/tools/pd-backup/pdbackup" + "github.com/pingcap/errors" + "github.com/pingcap/log" + "github.com/tikv/pd/tools/pd-dev/pd-backup/pdbackup" "go.etcd.io/etcd/clientv3" "go.etcd.io/etcd/pkg/transport" -) - -var ( - pdAddr = flag.String("pd", "http://127.0.0.1:2379", "pd address") - filePath = flag.String("file", "backup.json", "backup file path and name") - caPath = flag.String("cacert", "", "path of file that contains list of trusted SSL CAs") - certPath = flag.String("cert", "", "path of file that contains X509 certificate in PEM format") - keyPath = flag.String("key", "", "path of file that contains X509 key in PEM format") + "go.uber.org/zap" ) const ( etcdTimeout = 3 * time.Second ) -func main() { - flag.Parse() - f, err := os.Create(*filePath) +func Run() { + cfg := newConfig() + err := cfg.Parse(os.Args[1:]) + switch errors.Cause(err) { + case nil: + case flag.ErrHelp: + os.Exit(0) + default: + log.Fatal("parse cmd flags error", zap.Error(err)) + } + + f, err := os.Create(cfg.FilePath) checkErr(err) defer func() { if err := f.Close(); err != nil { @@ -48,12 +52,12 @@ func main() { } }() - urls := strings.Split(*pdAddr, ",") + urls := strings.Split(cfg.PDAddrs, ",") tlsInfo := transport.TLSInfo{ - CertFile: *certPath, - KeyFile: *keyPath, - TrustedCAFile: *caPath, + CertFile: cfg.CertPath, + KeyFile: cfg.KeyPath, + TrustedCAFile: cfg.CaPath, } tlsConfig, err := tlsInfo.ClientConfig() checkErr(err) @@ -65,10 +69,10 @@ func main() { }) checkErr(err) - backInfo, err := pdbackup.GetBackupInfo(client, *pdAddr) + backInfo, err := pdbackup.GetBackupInfo(client, cfg.PDAddrs) checkErr(err) pdbackup.OutputToFile(backInfo, f) - fmt.Println("pd backup successful! dump file is:", *filePath) + fmt.Println("pd backup successful! dump file is:", cfg.FilePath) } func checkErr(err error) { diff --git a/tools/pd-backup/pdbackup/backup.go b/tools/pd-dev/pd-backup/pdbackup/backup.go similarity index 100% rename from tools/pd-backup/pdbackup/backup.go rename to tools/pd-dev/pd-backup/pdbackup/backup.go diff --git a/tools/pd-backup/pdbackup/backup_test.go b/tools/pd-dev/pd-backup/pdbackup/backup_test.go similarity index 100% rename from tools/pd-backup/pdbackup/backup_test.go rename to tools/pd-dev/pd-backup/pdbackup/backup_test.go diff --git a/tools/pd-backup/tests/backup_test.go b/tools/pd-dev/pd-backup/tests/backup_test.go similarity index 96% rename from tools/pd-backup/tests/backup_test.go rename to tools/pd-dev/pd-backup/tests/backup_test.go index d7b7e613be9..69655a2a13a 100644 --- a/tools/pd-backup/tests/backup_test.go +++ b/tools/pd-dev/pd-backup/tests/backup_test.go @@ -24,7 +24,7 @@ import ( "github.com/stretchr/testify/require" "github.com/tikv/pd/tests" - "github.com/tikv/pd/tools/pd-backup/pdbackup" + "github.com/tikv/pd/tools/pd-dev/pd-backup/pdbackup" "go.etcd.io/etcd/clientv3" ) diff --git a/tools/pd-heartbeat-bench/config-template.toml b/tools/pd-dev/pd-heartbeat-bench/config-template.toml similarity index 100% rename from tools/pd-heartbeat-bench/config-template.toml rename to tools/pd-dev/pd-heartbeat-bench/config-template.toml diff --git a/tools/pd-heartbeat-bench/config/config.go b/tools/pd-dev/pd-heartbeat-bench/config.go similarity index 60% rename from tools/pd-heartbeat-bench/config/config.go rename to tools/pd-dev/pd-heartbeat-bench/config.go index 12455d78658..73d093322c2 100644 --- a/tools/pd-heartbeat-bench/config/config.go +++ b/tools/pd-dev/pd-heartbeat-bench/config.go @@ -1,14 +1,27 @@ -package config +// Copyright 2024 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 heartbeatbench import ( "sync/atomic" "github.com/BurntSushi/toml" "github.com/pingcap/errors" - "github.com/pingcap/log" flag "github.com/spf13/pflag" "github.com/tikv/pd/pkg/utils/configutil" - "go.uber.org/zap" + "github.com/tikv/pd/tools/pd-dev/util" ) const ( @@ -24,22 +37,11 @@ const ( defaultRound = 0 defaultSample = false defaultInitialVersion = 1 - - defaultLogFormat = "text" ) -// Config is the heartbeat-bench configuration. -type Config struct { - flagSet *flag.FlagSet - configFile string - PDAddr string - StatusAddr string - - Log log.Config `toml:"log" json:"log"` - Logger *zap.Logger - LogProps *log.ZapProperties - - Security configutil.SecurityConfig `toml:"security" json:"security"` +// config is the heartbeat-bench configuration. +type config struct { + *util.GeneralConfig InitEpochVer uint64 `toml:"epoch-ver" json:"epoch-ver"` StoreCount int `toml:"store-count" json:"store-count"` @@ -55,49 +57,45 @@ type Config struct { Round int `toml:"round" json:"round"` } -// NewConfig return a set of settings. -func NewConfig() *Config { - cfg := &Config{} - cfg.flagSet = flag.NewFlagSet("heartbeat-bench", flag.ContinueOnError) - fs := cfg.flagSet +// newConfig return a set of settings. +func newConfig() *config { + cfg := &config{ + GeneralConfig: &util.GeneralConfig{}, + } + cfg.FlagSet = flag.NewFlagSet("heartbeat-bench", flag.ContinueOnError) + fs := cfg.FlagSet + cfg.GeneralConfig = util.NewGeneralConfig(fs) fs.ParseErrorsWhitelist.UnknownFlags = true - fs.StringVar(&cfg.configFile, "config", "", "config file") - fs.StringVar(&cfg.PDAddr, "pd-endpoints", "127.0.0.1:2379", "pd address") - fs.StringVar(&cfg.Log.File.Filename, "log-file", "", "log file path") - fs.StringVar(&cfg.StatusAddr, "status-addr", "127.0.0.1:20180", "status address") - fs.StringVar(&cfg.Security.CAPath, "cacert", "", "path of file that contains list of trusted TLS CAs") - fs.StringVar(&cfg.Security.CertPath, "cert", "", "path of file that contains X509 certificate in PEM format") - fs.StringVar(&cfg.Security.KeyPath, "key", "", "path of file that contains X509 key in PEM format") fs.Uint64Var(&cfg.InitEpochVer, "epoch-ver", 1, "the initial epoch version value") return cfg } // Parse parses flag definitions from the argument list. -func (c *Config) Parse(arguments []string) error { +func (c *config) Parse(arguments []string) error { // Parse first to get config file. - err := c.flagSet.Parse(arguments) + err := c.FlagSet.Parse(arguments) if err != nil { return errors.WithStack(err) } // Load config file if specified. var meta *toml.MetaData - if c.configFile != "" { - meta, err = configutil.ConfigFromFile(c, c.configFile) + if c.ConfigFile != "" { + meta, err = configutil.ConfigFromFile(c, c.ConfigFile) if err != nil { return err } } // Parse again to replace with command line options. - err = c.flagSet.Parse(arguments) + err = c.FlagSet.Parse(arguments) if err != nil { return errors.WithStack(err) } - if len(c.flagSet.Args()) != 0 { - return errors.Errorf("'%s' is an invalid flag", c.flagSet.Arg(0)) + if len(c.FlagSet.Args()) != 0 { + return errors.Errorf("'%s' is an invalid flag", c.FlagSet.Arg(0)) } c.Adjust(meta) @@ -105,10 +103,9 @@ func (c *Config) Parse(arguments []string) error { } // Adjust is used to adjust configurations -func (c *Config) Adjust(meta *toml.MetaData) { - if len(c.Log.Format) == 0 { - c.Log.Format = defaultLogFormat - } +func (c *config) Adjust(meta *toml.MetaData) { + c.GeneralConfig.Adjust() + if !meta.IsDefined("round") { configutil.AdjustInt(&c.Round, defaultRound) } @@ -151,7 +148,7 @@ func (c *Config) Adjust(meta *toml.MetaData) { } // Validate is used to validate configurations -func (c *Config) Validate() error { +func (c *config) Validate() error { if c.HotStoreCount < 0 || c.HotStoreCount > c.StoreCount { return errors.Errorf("hot-store-count must be in [0, store-count]") } @@ -174,71 +171,71 @@ func (c *Config) Validate() error { } // Clone creates a copy of current config. -func (c *Config) Clone() *Config { - cfg := &Config{} +func (c *config) Clone() *config { + cfg := &config{} *cfg = *c return cfg } // Options is the option of the heartbeat-bench. -type Options struct { - HotStoreCount atomic.Value - ReportRatio atomic.Value - - LeaderUpdateRatio atomic.Value - EpochUpdateRatio atomic.Value - SpaceUpdateRatio atomic.Value - FlowUpdateRatio atomic.Value +type options struct { + hotStoreCount atomic.Value + reportRatio atomic.Value + + leaderUpdateRatio atomic.Value + epochUpdateRatio atomic.Value + spaceUpdateRatio atomic.Value + flowUpdateRatio atomic.Value } // NewOptions creates a new option. -func NewOptions(cfg *Config) *Options { - o := &Options{} - o.HotStoreCount.Store(cfg.HotStoreCount) - o.LeaderUpdateRatio.Store(cfg.LeaderUpdateRatio) - o.EpochUpdateRatio.Store(cfg.EpochUpdateRatio) - o.SpaceUpdateRatio.Store(cfg.SpaceUpdateRatio) - o.FlowUpdateRatio.Store(cfg.FlowUpdateRatio) - o.ReportRatio.Store(cfg.ReportRatio) +func newOptions(cfg *config) *options { + o := &options{} + o.hotStoreCount.Store(cfg.HotStoreCount) + o.leaderUpdateRatio.Store(cfg.LeaderUpdateRatio) + o.epochUpdateRatio.Store(cfg.EpochUpdateRatio) + o.spaceUpdateRatio.Store(cfg.SpaceUpdateRatio) + o.flowUpdateRatio.Store(cfg.FlowUpdateRatio) + o.reportRatio.Store(cfg.ReportRatio) return o } // GetHotStoreCount returns the hot store count. -func (o *Options) GetHotStoreCount() int { - return o.HotStoreCount.Load().(int) +func (o *options) GetHotStoreCount() int { + return o.hotStoreCount.Load().(int) } // GetLeaderUpdateRatio returns the leader update ratio. -func (o *Options) GetLeaderUpdateRatio() float64 { - return o.LeaderUpdateRatio.Load().(float64) +func (o *options) GetLeaderUpdateRatio() float64 { + return o.leaderUpdateRatio.Load().(float64) } // GetEpochUpdateRatio returns the epoch update ratio. -func (o *Options) GetEpochUpdateRatio() float64 { - return o.EpochUpdateRatio.Load().(float64) +func (o *options) GetEpochUpdateRatio() float64 { + return o.epochUpdateRatio.Load().(float64) } // GetSpaceUpdateRatio returns the space update ratio. -func (o *Options) GetSpaceUpdateRatio() float64 { - return o.SpaceUpdateRatio.Load().(float64) +func (o *options) GetSpaceUpdateRatio() float64 { + return o.spaceUpdateRatio.Load().(float64) } // GetFlowUpdateRatio returns the flow update ratio. -func (o *Options) GetFlowUpdateRatio() float64 { - return o.FlowUpdateRatio.Load().(float64) +func (o *options) GetFlowUpdateRatio() float64 { + return o.flowUpdateRatio.Load().(float64) } // GetReportRatio returns the report ratio. -func (o *Options) GetReportRatio() float64 { - return o.ReportRatio.Load().(float64) +func (o *options) GetReportRatio() float64 { + return o.reportRatio.Load().(float64) } // SetOptions sets the option. -func (o *Options) SetOptions(cfg *Config) { - o.HotStoreCount.Store(cfg.HotStoreCount) - o.LeaderUpdateRatio.Store(cfg.LeaderUpdateRatio) - o.EpochUpdateRatio.Store(cfg.EpochUpdateRatio) - o.SpaceUpdateRatio.Store(cfg.SpaceUpdateRatio) - o.FlowUpdateRatio.Store(cfg.FlowUpdateRatio) - o.ReportRatio.Store(cfg.ReportRatio) +func (o *options) SetOptions(cfg *config) { + o.hotStoreCount.Store(cfg.HotStoreCount) + o.leaderUpdateRatio.Store(cfg.LeaderUpdateRatio) + o.epochUpdateRatio.Store(cfg.EpochUpdateRatio) + o.spaceUpdateRatio.Store(cfg.SpaceUpdateRatio) + o.flowUpdateRatio.Store(cfg.FlowUpdateRatio) + o.reportRatio.Store(cfg.ReportRatio) } diff --git a/tools/pd-dev/pd-heartbeat-bench/main.go b/tools/pd-dev/pd-heartbeat-bench/main.go new file mode 100644 index 00000000000..7dc58ca568e --- /dev/null +++ b/tools/pd-dev/pd-heartbeat-bench/main.go @@ -0,0 +1,202 @@ +// Copyright 2019 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 heartbeatbench + +import ( + "context" + "fmt" + "math/rand" + "net/http" + "os" + "sync" + "time" + + "github.com/gin-contrib/cors" + "github.com/gin-contrib/gzip" + "github.com/gin-contrib/pprof" + "github.com/gin-gonic/gin" + "github.com/pingcap/errors" + "github.com/pingcap/kvproto/pkg/pdpb" + "github.com/pingcap/log" + "github.com/spf13/pflag" + pdHttp "github.com/tikv/pd/client/http" + "github.com/tikv/pd/pkg/mcs/utils" + "github.com/tikv/pd/pkg/statistics" + "github.com/tikv/pd/pkg/utils/logutil" + "github.com/tikv/pd/tools/pd-dev/util" + "go.uber.org/zap" +) + +func Run(ctx context.Context) { + defer logutil.LogPanic() + + rand.New(rand.NewSource(0)) // Ensure consistent behavior multiple times + statistics.Denoising = false + + cfg := newConfig() + err := cfg.Parse(os.Args[1:]) + switch errors.Cause(err) { + case nil: + case pflag.ErrHelp: + os.Exit(0) + default: + log.Fatal("parse cmd flags error", zap.Error(err)) + } + + maxVersion = cfg.InitEpochVer + options := newOptions(cfg) + // let PD have enough time to start + time.Sleep(5 * time.Second) + + cli, err := newClient(ctx, cfg) + if err != nil { + log.Fatal("create client error", zap.Error(err)) + } + + initClusterID(ctx, cli) + srv := runHTTPServer(cfg, options) + regions := new(Regions) + regions.init(cfg) + log.Info("finish init regions") + stores := newStores(cfg.StoreCount) + stores.update(regions) + bootstrap(ctx, cli) + putStores(ctx, cfg, cli, stores) + log.Info("finish put stores") + clis := make(map[uint64]pdpb.PDClient, cfg.StoreCount) + httpCli := pdHttp.NewClient("tools-heartbeat-bench", []string{cfg.PDAddrs}, pdHttp.WithTLSConfig(util.LoadTLSConfig(cfg.GeneralConfig))) + go deleteOperators(ctx, httpCli) + streams := make(map[uint64]pdpb.PD_RegionHeartbeatClient, cfg.StoreCount) + for i := 1; i <= cfg.StoreCount; i++ { + clis[uint64(i)], streams[uint64(i)] = createHeartbeatStream(ctx, cfg) + } + header := &pdpb.RequestHeader{ + ClusterId: clusterID, + } + var heartbeatTicker = time.NewTicker(regionReportInterval * time.Second) + defer heartbeatTicker.Stop() + var resolvedTSTicker = time.NewTicker(time.Second) + defer resolvedTSTicker.Stop() + for { + select { + case <-heartbeatTicker.C: + if cfg.Round != 0 && regions.updateRound > cfg.Round { + os.Exit(0) + } + rep := newReport(cfg) + r := rep.Stats() + + startTime := time.Now() + wg := &sync.WaitGroup{} + for i := 1; i <= cfg.StoreCount; i++ { + id := uint64(i) + wg.Add(1) + go regions.handleRegionHeartbeat(wg, streams[id], id, rep) + } + wg.Wait() + + since := time.Since(startTime).Seconds() + close(rep.Results()) + regions.result(cfg.RegionCount, since) + stats := <-r + log.Info("region heartbeat stats", zap.String("total", fmt.Sprintf("%.4fs", stats.Total.Seconds())), + zap.String("slowest", fmt.Sprintf("%.4fs", stats.Slowest)), + zap.String("fastest", fmt.Sprintf("%.4fs", stats.Fastest)), + zap.String("average", fmt.Sprintf("%.4fs", stats.Average)), + zap.String("stddev", fmt.Sprintf("%.4fs", stats.Stddev)), + zap.String("rps", fmt.Sprintf("%.4f", stats.RPS)), + zap.Uint64("max-epoch-version", maxVersion), + ) + log.Info("store heartbeat stats", zap.String("max", fmt.Sprintf("%.4fs", since))) + regions.update(cfg, options) + go stores.update(regions) // update stores in background, unusually region heartbeat is slower than store update. + case <-resolvedTSTicker.C: + wg := &sync.WaitGroup{} + for i := 1; i <= cfg.StoreCount; i++ { + id := uint64(i) + wg.Add(1) + go func(wg *sync.WaitGroup, id uint64) { + defer wg.Done() + cli := clis[id] + _, err := cli.ReportMinResolvedTS(ctx, &pdpb.ReportMinResolvedTsRequest{ + Header: header, + StoreId: id, + MinResolvedTs: uint64(time.Now().Unix()), + }) + if err != nil { + log.Error("send resolved TS error", zap.Uint64("store-id", id), zap.Error(err)) + return + } + }(wg, id) + } + wg.Wait() + case <-ctx.Done(): + if err := srv.Shutdown(context.Background()); err != nil { + log.Fatal("server shutdown error", zap.Error(err)) + } + log.Info("got signal to exit") + return + } + } +} + +func runHTTPServer(cfg *config, options *options) *http.Server { + gin.SetMode(gin.ReleaseMode) + engine := gin.New() + engine.Use(gin.Recovery()) + engine.Use(cors.Default()) + engine.Use(gzip.Gzip(gzip.DefaultCompression)) + engine.GET("metrics", utils.PromHandler()) + // profile API + pprof.Register(engine) + engine.PUT("config", func(c *gin.Context) { + newCfg := cfg.Clone() + newCfg.HotStoreCount = options.GetHotStoreCount() + newCfg.FlowUpdateRatio = options.GetFlowUpdateRatio() + newCfg.LeaderUpdateRatio = options.GetLeaderUpdateRatio() + newCfg.EpochUpdateRatio = options.GetEpochUpdateRatio() + newCfg.SpaceUpdateRatio = options.GetSpaceUpdateRatio() + newCfg.ReportRatio = options.GetReportRatio() + if err := c.BindJSON(&newCfg); err != nil { + c.String(http.StatusBadRequest, err.Error()) + return + } + if err := newCfg.Validate(); err != nil { + c.String(http.StatusBadRequest, err.Error()) + return + } + options.SetOptions(newCfg) + c.String(http.StatusOK, "Successfully updated the configuration") + }) + engine.GET("config", func(c *gin.Context) { + output := cfg.Clone() + output.HotStoreCount = options.GetHotStoreCount() + output.FlowUpdateRatio = options.GetFlowUpdateRatio() + output.LeaderUpdateRatio = options.GetLeaderUpdateRatio() + output.EpochUpdateRatio = options.GetEpochUpdateRatio() + output.SpaceUpdateRatio = options.GetSpaceUpdateRatio() + output.ReportRatio = options.GetReportRatio() + + c.IndentedJSON(http.StatusOK, output) + }) + srv := &http.Server{Addr: cfg.StatusAddr, Handler: engine.Handler(), ReadHeaderTimeout: 3 * time.Second} + go func() { + if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed { + log.Fatal("server listen error", zap.Error(err)) + } + }() + + return srv +} diff --git a/tools/pd-dev/pd-heartbeat-bench/region.go b/tools/pd-dev/pd-heartbeat-bench/region.go new file mode 100644 index 00000000000..033427b0c58 --- /dev/null +++ b/tools/pd-dev/pd-heartbeat-bench/region.go @@ -0,0 +1,260 @@ +// Copyright 2024 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 heartbeatbench + +import ( + "context" + "fmt" + "io" + "math/rand" + "sync" + "sync/atomic" + "time" + + "github.com/pingcap/kvproto/pkg/metapb" + "github.com/pingcap/kvproto/pkg/pdpb" + "github.com/pingcap/log" + "github.com/tikv/pd/pkg/codec" + "go.etcd.io/etcd/pkg/report" + "go.uber.org/zap" +) + +// Regions simulates all regions to heartbeat. +type Regions struct { + regions []*pdpb.RegionHeartbeatRequest + awakenRegions atomic.Value + + updateRound int + + updateLeader []int + updateEpoch []int + updateSpace []int + updateFlow []int +} + +func (rs *Regions) init(cfg *config) { + rs.regions = make([]*pdpb.RegionHeartbeatRequest, 0, cfg.RegionCount) + rs.updateRound = 0 + + // Generate regions + id := uint64(1) + now := uint64(time.Now().Unix()) + + for i := 0; i < cfg.RegionCount; i++ { + region := &pdpb.RegionHeartbeatRequest{ + Header: header(), + Region: &metapb.Region{ + Id: id, + StartKey: codec.GenerateTableKey(int64(i)), + EndKey: codec.GenerateTableKey(int64(i + 1)), + RegionEpoch: &metapb.RegionEpoch{ConfVer: 2, Version: maxVersion}, + }, + ApproximateSize: bytesUnit, + Interval: &pdpb.TimeInterval{ + StartTimestamp: now, + EndTimestamp: now + regionReportInterval, + }, + QueryStats: &pdpb.QueryStats{}, + ApproximateKeys: keysUint, + Term: 1, + } + id += 1 + if i == 0 { + region.Region.StartKey = []byte("") + } + if i == cfg.RegionCount-1 { + region.Region.EndKey = []byte("") + } + + peers := make([]*metapb.Peer, 0, cfg.Replica) + for j := 0; j < cfg.Replica; j++ { + peers = append(peers, &metapb.Peer{Id: id, StoreId: uint64((i+j)%cfg.StoreCount + 1)}) + id += 1 + } + + region.Region.Peers = peers + region.Leader = peers[0] + rs.regions = append(rs.regions, region) + } +} + +func (rs *Regions) update(cfg *config, options *options) { + rs.updateRound += 1 + + // Generate sample index + indexes := make([]int, cfg.RegionCount) + for i := range indexes { + indexes[i] = i + } + reportRegions := pick(indexes, cfg.RegionCount, options.GetReportRatio()) + + reportCount := len(reportRegions) + rs.updateFlow = pick(reportRegions, reportCount, options.GetFlowUpdateRatio()) + rs.updateLeader = randomPick(reportRegions, reportCount, options.GetLeaderUpdateRatio()) + rs.updateEpoch = randomPick(reportRegions, reportCount, options.GetEpochUpdateRatio()) + rs.updateSpace = randomPick(reportRegions, reportCount, options.GetSpaceUpdateRatio()) + var ( + updatedStatisticsMap = make(map[int]*pdpb.RegionHeartbeatRequest) + awakenRegions []*pdpb.RegionHeartbeatRequest + ) + + // update leader + for _, i := range rs.updateLeader { + region := rs.regions[i] + region.Leader = region.Region.Peers[rs.updateRound%cfg.Replica] + } + // update epoch + for _, i := range rs.updateEpoch { + region := rs.regions[i] + region.Region.RegionEpoch.Version += 1 + if region.Region.RegionEpoch.Version > maxVersion { + maxVersion = region.Region.RegionEpoch.Version + } + } + // update space + for _, i := range rs.updateSpace { + region := rs.regions[i] + region.ApproximateSize = uint64(bytesUnit * rand.Float64()) + region.ApproximateKeys = uint64(keysUint * rand.Float64()) + } + // update flow + for _, i := range rs.updateFlow { + region := rs.regions[i] + if region.Leader.StoreId <= uint64(options.GetHotStoreCount()) { + region.BytesWritten = uint64(hotByteUnit * (1 + rand.Float64()) * 60) + region.BytesRead = uint64(hotByteUnit * (1 + rand.Float64()) * 10) + region.KeysWritten = uint64(hotKeysUint * (1 + rand.Float64()) * 60) + region.KeysRead = uint64(hotKeysUint * (1 + rand.Float64()) * 10) + region.QueryStats = &pdpb.QueryStats{ + Get: uint64(hotQueryUnit * (1 + rand.Float64()) * 10), + Put: uint64(hotQueryUnit * (1 + rand.Float64()) * 60), + } + } else { + region.BytesWritten = uint64(bytesUnit * rand.Float64()) + region.BytesRead = uint64(bytesUnit * rand.Float64()) + region.KeysWritten = uint64(keysUint * rand.Float64()) + region.KeysRead = uint64(keysUint * rand.Float64()) + region.QueryStats = &pdpb.QueryStats{ + Get: uint64(queryUnit * rand.Float64()), + Put: uint64(queryUnit * rand.Float64()), + } + } + updatedStatisticsMap[i] = region + } + // update interval + for _, region := range rs.regions { + region.Interval.StartTimestamp = region.Interval.EndTimestamp + region.Interval.EndTimestamp = region.Interval.StartTimestamp + regionReportInterval + } + for _, i := range reportRegions { + region := rs.regions[i] + // reset the statistics of the region which is not updated + if _, exist := updatedStatisticsMap[i]; !exist { + region.BytesWritten = 0 + region.BytesRead = 0 + region.KeysWritten = 0 + region.KeysRead = 0 + region.QueryStats = &pdpb.QueryStats{} + } + awakenRegions = append(awakenRegions, region) + } + + rs.awakenRegions.Store(awakenRegions) +} + +func createHeartbeatStream(ctx context.Context, cfg *config) (pdpb.PDClient, pdpb.PD_RegionHeartbeatClient) { + cli, err := newClient(ctx, cfg) + if err != nil { + log.Fatal("create client error", zap.Error(err)) + } + stream, err := cli.RegionHeartbeat(ctx) + if err != nil { + log.Fatal("create stream error", zap.Error(err)) + } + + go func() { + // do nothing + for { + stream.Recv() + } + }() + return cli, stream +} + +func (rs *Regions) handleRegionHeartbeat(wg *sync.WaitGroup, stream pdpb.PD_RegionHeartbeatClient, storeID uint64, rep report.Report) { + defer wg.Done() + var regions, toUpdate []*pdpb.RegionHeartbeatRequest + updatedRegions := rs.awakenRegions.Load() + if updatedRegions == nil { + toUpdate = rs.regions + } else { + toUpdate = updatedRegions.([]*pdpb.RegionHeartbeatRequest) + } + for _, region := range toUpdate { + if region.Leader.StoreId != storeID { + continue + } + regions = append(regions, region) + } + + start := time.Now() + var err error + for _, region := range regions { + err = stream.Send(region) + rep.Results() <- report.Result{Start: start, End: time.Now(), Err: err} + if err == io.EOF { + log.Error("receive eof error", zap.Uint64("store-id", storeID), zap.Error(err)) + err := stream.CloseSend() + if err != nil { + log.Error("fail to close stream", zap.Uint64("store-id", storeID), zap.Error(err)) + } + return + } + if err != nil { + log.Error("send result error", zap.Uint64("store-id", storeID), zap.Error(err)) + return + } + } + log.Info("store finish one round region heartbeat", zap.Uint64("store-id", storeID), zap.Duration("cost-time", time.Since(start)), zap.Int("reported-region-count", len(regions))) +} + +func (rs *Regions) result(regionCount int, sec float64) { + if rs.updateRound == 0 { + // There was no difference in the first round + return + } + + updated := make(map[int]struct{}) + for _, i := range rs.updateLeader { + updated[i] = struct{}{} + } + for _, i := range rs.updateEpoch { + updated[i] = struct{}{} + } + for _, i := range rs.updateSpace { + updated[i] = struct{}{} + } + for _, i := range rs.updateFlow { + updated[i] = struct{}{} + } + inactiveCount := regionCount - len(updated) + + log.Info("update speed of each category", zap.String("rps", fmt.Sprintf("%.4f", float64(regionCount)/sec)), + zap.String("save-tree", fmt.Sprintf("%.4f", float64(len(rs.updateLeader))/sec)), + zap.String("save-kv", fmt.Sprintf("%.4f", float64(len(rs.updateEpoch))/sec)), + zap.String("save-space", fmt.Sprintf("%.4f", float64(len(rs.updateSpace))/sec)), + zap.String("save-flow", fmt.Sprintf("%.4f", float64(len(rs.updateFlow))/sec)), + zap.String("skip", fmt.Sprintf("%.4f", float64(inactiveCount)/sec))) +} diff --git a/tools/pd-dev/pd-heartbeat-bench/store.go b/tools/pd-dev/pd-heartbeat-bench/store.go new file mode 100644 index 00000000000..d4747814c21 --- /dev/null +++ b/tools/pd-dev/pd-heartbeat-bench/store.go @@ -0,0 +1,93 @@ +// Copyright 2024 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 heartbeatbench + +import ( + "context" + "sync/atomic" + "time" + + "github.com/pingcap/kvproto/pkg/pdpb" +) + +// Stores contains store stats with lock. +type Stores struct { + stat []atomic.Value +} + +func newStores(storeCount int) *Stores { + return &Stores{ + stat: make([]atomic.Value, storeCount+1), + } +} + +func (s *Stores) heartbeat(ctx context.Context, cli pdpb.PDClient, storeID uint64) { + cctx, cancel := context.WithCancel(ctx) + defer cancel() + cli.StoreHeartbeat(cctx, &pdpb.StoreHeartbeatRequest{Header: header(), Stats: s.stat[storeID].Load().(*pdpb.StoreStats)}) +} + +func (s *Stores) update(rs *Regions) { + stats := make([]*pdpb.StoreStats, len(s.stat)) + now := uint64(time.Now().Unix()) + for i := range stats { + stats[i] = &pdpb.StoreStats{ + StoreId: uint64(i), + Capacity: capacity, + Available: capacity, + QueryStats: &pdpb.QueryStats{}, + PeerStats: make([]*pdpb.PeerStat, 0), + Interval: &pdpb.TimeInterval{ + StartTimestamp: now - storeReportInterval, + EndTimestamp: now, + }, + } + } + var toUpdate []*pdpb.RegionHeartbeatRequest + updatedRegions := rs.awakenRegions.Load() + if updatedRegions == nil { + toUpdate = rs.regions + } else { + toUpdate = updatedRegions.([]*pdpb.RegionHeartbeatRequest) + } + for _, region := range toUpdate { + for _, peer := range region.Region.Peers { + store := stats[peer.StoreId] + store.UsedSize += region.ApproximateSize + store.Available -= region.ApproximateSize + store.RegionCount += 1 + } + store := stats[region.Leader.StoreId] + if region.BytesWritten != 0 { + store.BytesWritten += region.BytesWritten + store.BytesRead += region.BytesRead + store.KeysWritten += region.KeysWritten + store.KeysRead += region.KeysRead + store.QueryStats.Get += region.QueryStats.Get + store.QueryStats.Put += region.QueryStats.Put + store.PeerStats = append(store.PeerStats, &pdpb.PeerStat{ + RegionId: region.Region.Id, + ReadKeys: region.KeysRead, + ReadBytes: region.BytesRead, + WrittenKeys: region.KeysWritten, + WrittenBytes: region.BytesWritten, + QueryStats: region.QueryStats, + }) + } + } + for i := range stats { + s.stat[i].Store(stats[i]) + } +} diff --git a/tools/pd-dev/pd-heartbeat-bench/util.go b/tools/pd-dev/pd-heartbeat-bench/util.go new file mode 100644 index 00000000000..b43ba920234 --- /dev/null +++ b/tools/pd-dev/pd-heartbeat-bench/util.go @@ -0,0 +1,193 @@ +// Copyright 2024 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 heartbeatbench + +import ( + "context" + "fmt" + "math/rand" + "time" + + "github.com/docker/go-units" + "github.com/pingcap/kvproto/pkg/metapb" + "github.com/pingcap/kvproto/pkg/pdpb" + "github.com/pingcap/log" + pdHttp "github.com/tikv/pd/client/http" + "github.com/tikv/pd/pkg/utils/grpcutil" + "github.com/tikv/pd/tools/pd-dev/util" + "go.etcd.io/etcd/pkg/report" + "go.uber.org/zap" +) + +const ( + bytesUnit = 128 + keysUint = 8 + queryUnit = 8 + hotByteUnit = 16 * units.KiB + hotKeysUint = 256 + hotQueryUnit = 256 + regionReportInterval = 60 // 60s + storeReportInterval = 10 // 10s + capacity = 4 * units.TiB +) + +var ( + clusterID uint64 + maxVersion uint64 = 1 +) + +func newClient(ctx context.Context, cfg *config) (pdpb.PDClient, error) { + tlsConfig := util.LoadTLSConfig(cfg.GeneralConfig) + cc, err := grpcutil.GetClientConn(ctx, cfg.PDAddrs, tlsConfig) + if err != nil { + return nil, err + } + return pdpb.NewPDClient(cc), nil +} + +func initClusterID(ctx context.Context, cli pdpb.PDClient) { + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + cctx, cancel := context.WithCancel(ctx) + res, err := cli.GetMembers(cctx, &pdpb.GetMembersRequest{}) + cancel() + if err != nil { + continue + } + if res.GetHeader().GetError() != nil { + continue + } + clusterID = res.GetHeader().GetClusterId() + log.Info("init cluster ID successfully", zap.Uint64("cluster-id", clusterID)) + return + } + } +} + +func header() *pdpb.RequestHeader { + return &pdpb.RequestHeader{ + ClusterId: clusterID, + } +} + +func bootstrap(ctx context.Context, cli pdpb.PDClient) { + cctx, cancel := context.WithCancel(ctx) + isBootstrapped, err := cli.IsBootstrapped(cctx, &pdpb.IsBootstrappedRequest{Header: header()}) + cancel() + if err != nil { + log.Fatal("check if cluster has already bootstrapped failed", zap.Error(err)) + } + if isBootstrapped.GetBootstrapped() { + log.Info("already bootstrapped") + return + } + + store := &metapb.Store{ + Id: 1, + Address: fmt.Sprintf("localhost:%d", 2), + Version: "8.0.0-alpha", + } + region := &metapb.Region{ + Id: 1, + Peers: []*metapb.Peer{{StoreId: 1, Id: 1}}, + RegionEpoch: &metapb.RegionEpoch{ConfVer: 1, Version: 1}, + } + req := &pdpb.BootstrapRequest{ + Header: header(), + Store: store, + Region: region, + } + cctx, cancel = context.WithCancel(ctx) + resp, err := cli.Bootstrap(cctx, req) + cancel() + if err != nil { + log.Fatal("failed to bootstrap the cluster", zap.Error(err)) + } + if resp.GetHeader().GetError() != nil { + log.Fatal("failed to bootstrap the cluster", zap.String("err", resp.GetHeader().GetError().String())) + } + log.Info("bootstrapped") +} + +func putStores(ctx context.Context, cfg *config, cli pdpb.PDClient, stores *Stores) { + for i := uint64(1); i <= uint64(cfg.StoreCount); i++ { + store := &metapb.Store{ + Id: i, + Address: fmt.Sprintf("localhost:%d", i), + Version: "8.0.0-alpha", + } + cctx, cancel := context.WithCancel(ctx) + resp, err := cli.PutStore(cctx, &pdpb.PutStoreRequest{Header: header(), Store: store}) + cancel() + if err != nil { + log.Fatal("failed to put store", zap.Uint64("store-id", i), zap.Error(err)) + } + if resp.GetHeader().GetError() != nil { + log.Fatal("failed to put store", zap.Uint64("store-id", i), zap.String("err", resp.GetHeader().GetError().String())) + } + go func(ctx context.Context, storeID uint64) { + var heartbeatTicker = time.NewTicker(10 * time.Second) + defer heartbeatTicker.Stop() + for { + select { + case <-heartbeatTicker.C: + stores.heartbeat(ctx, cli, storeID) + case <-ctx.Done(): + return + } + } + }(ctx, i) + } +} + +func deleteOperators(ctx context.Context, httpCli pdHttp.Client) { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + err := httpCli.DeleteOperators(ctx) + if err != nil { + log.Error("fail to delete operators", zap.Error(err)) + } + } + } +} + +func newReport(cfg *config) report.Report { + p := "%4.4f" + if cfg.Sample { + return report.NewReportSample(p) + } + return report.NewReport(p) +} + +func randomPick(slice []int, total int, ratio float64) []int { + rand.Shuffle(total, func(i, j int) { + slice[i], slice[j] = slice[j], slice[i] + }) + return append(slice[:0:0], slice[0:int(float64(total)*ratio)]...) +} + +func pick(slice []int, total int, ratio float64) []int { + return append(slice[:0:0], slice[0:int(float64(total)*ratio)]...) +} diff --git a/tools/pd-simulator/README.md b/tools/pd-dev/pd-simulator/README.md similarity index 100% rename from tools/pd-simulator/README.md rename to tools/pd-dev/pd-simulator/README.md diff --git a/tools/pd-simulator/main.go b/tools/pd-dev/pd-simulator/main.go similarity index 76% rename from tools/pd-simulator/main.go rename to tools/pd-dev/pd-simulator/main.go index 73f4a0bba12..a26f4704510 100644 --- a/tools/pd-simulator/main.go +++ b/tools/pd-dev/pd-simulator/main.go @@ -12,16 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package simulator import ( "context" + "errors" "fmt" "net/http" "net/http/pprof" "os" - "os/signal" - "syscall" "time" "github.com/BurntSushi/toml" @@ -34,11 +33,12 @@ import ( "github.com/tikv/pd/pkg/utils/testutil" "github.com/tikv/pd/server" "github.com/tikv/pd/server/api" - "github.com/tikv/pd/server/config" - "github.com/tikv/pd/tools/pd-analysis/analysis" - "github.com/tikv/pd/tools/pd-simulator/simulator" - "github.com/tikv/pd/tools/pd-simulator/simulator/cases" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + svrCfg "github.com/tikv/pd/server/config" + "github.com/tikv/pd/tools/pd-dev/pd-analysis/analysis" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/cases" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/config" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) @@ -55,7 +55,7 @@ var ( statusAddress = flag.String("status-addr", "0.0.0.0:20180", "status address") ) -func main() { +func Run(ctx context.Context) { // wait PD start. Otherwise, it will happen error when getting cluster ID. time.Sleep(3 * time.Second) // ignore some undefined flag @@ -63,14 +63,14 @@ func main() { flag.Parse() simutil.InitLogger(*simLogLevel, *simLogFile) - simutil.InitCaseConfig(*storeNum, *regionNum, *enableTransferRegionCounter) + config.InitCaseConfig(*storeNum, *regionNum, *enableTransferRegionCounter) statistics.Denoising = false - if simutil.CaseConfigure.EnableTransferRegionCounter { - analysis.GetTransferCounter().Init(simutil.CaseConfigure.StoreNum, simutil.CaseConfigure.RegionNum) + if config.CaseConfigure.EnableTransferRegionCounter { + analysis.GetTransferCounter().Init(config.CaseConfigure.StoreNum, config.CaseConfigure.RegionNum) } schedulers.Register() // register schedulers, which is needed by simConfig.Adjust - simConfig := simulator.NewSimConfig(*serverLogLevel) + simConfig := config.NewSimConfig(*serverLogLevel) var meta toml.MetaData var err error if *configFile != "" { @@ -90,17 +90,20 @@ func main() { simutil.Logger.Fatal("need to specify one config name") } for simCase := range cases.CaseMap { - run(simCase, simConfig) + simRun(ctx, simCase, simConfig) } } else { - run(*caseName, simConfig) + simRun(ctx, *caseName, simConfig) } } -func run(simCase string, simConfig *simulator.SimConfig) { +func simRun(ctx context.Context, simCase string, simConfig *config.SimConfig) { if *pdAddr != "" { - go runHTTPServer() - simStart(*pdAddr, simCase, simConfig) + srv := runHTTPServer() + simStart(ctx, *pdAddr, simCase, simConfig) + if err := srv.Shutdown(context.Background()); err != nil { + log.Fatal("server shutdown error", zap.Error(err)) + } } else { local, clean := NewSingleServer(context.Background(), simConfig) err := local.Run() @@ -113,11 +116,11 @@ func run(simCase string, simConfig *simulator.SimConfig) { } time.Sleep(100 * time.Millisecond) } - simStart(local.GetAddr(), simCase, simConfig, clean) + simStart(ctx, local.GetAddr(), simCase, simConfig, clean) } } -func runHTTPServer() { +func runHTTPServer() *http.Server { http.Handle("/metrics", promhttp.Handler()) // profile API http.HandleFunc("/pprof/profile", pprof.Profile) @@ -128,15 +131,20 @@ func runHTTPServer() { http.Handle("/pprof/allocs", pprof.Handler("allocs")) http.Handle("/pprof/block", pprof.Handler("block")) http.Handle("/pprof/goroutine", pprof.Handler("goroutine")) - server := &http.Server{ + srv := &http.Server{ Addr: *statusAddress, ReadHeaderTimeout: 3 * time.Second, } - server.ListenAndServe() + go func() { + if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + log.Fatal("server listen error", zap.Error(err)) + } + }() + return srv } // NewSingleServer creates a pd server for simulator. -func NewSingleServer(ctx context.Context, simConfig *simulator.SimConfig) (*server.Server, testutil.CleanupFunc) { +func NewSingleServer(ctx context.Context, simConfig *config.SimConfig) (*server.Server, testutil.CleanupFunc) { err := logutil.SetupLogger(simConfig.ServerConfig.Log, &simConfig.ServerConfig.Logger, &simConfig.ServerConfig.LogProps) if err == nil { log.ReplaceGlobals(simConfig.ServerConfig.Logger, simConfig.ServerConfig.LogProps) @@ -156,12 +164,13 @@ func NewSingleServer(ctx context.Context, simConfig *simulator.SimConfig) (*serv return s, cleanup } -func cleanServer(cfg *config.Config) { +func cleanServer(cfg *svrCfg.Config) { // Clean data directory os.RemoveAll(cfg.DataDir) } -func simStart(pdAddr string, simCase string, simConfig *simulator.SimConfig, clean ...testutil.CleanupFunc) { +func simStart(ctx context.Context, pdAddr string, + simCase string, simConfig *config.SimConfig, clean ...testutil.CleanupFunc) { start := time.Now() driver, err := simulator.NewDriver(pdAddr, simCase, simConfig) if err != nil { @@ -176,12 +185,6 @@ func simStart(pdAddr string, simCase string, simConfig *simulator.SimConfig, cle tick := time.NewTicker(tickInterval) defer tick.Stop() - sc := make(chan os.Signal, 1) - signal.Notify(sc, - syscall.SIGHUP, - syscall.SIGINT, - syscall.SIGTERM, - syscall.SIGQUIT) simResult := "FAIL" @@ -194,7 +197,7 @@ EXIT: simResult = "OK" break EXIT } - case <-sc: + case <-ctx.Done(): break EXIT } } diff --git a/tools/pd-simulator/simulator/cases/add_nodes.go b/tools/pd-dev/pd-simulator/simulator/cases/add_nodes.go similarity index 94% rename from tools/pd-simulator/simulator/cases/add_nodes.go rename to tools/pd-dev/pd-simulator/simulator/cases/add_nodes.go index 241b34a9473..9e864eb71f9 100644 --- a/tools/pd-simulator/simulator/cases/add_nodes.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/add_nodes.go @@ -20,8 +20,8 @@ import ( "github.com/docker/go-units" "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/add_nodes_dynamic.go b/tools/pd-dev/pd-simulator/simulator/cases/add_nodes_dynamic.go similarity index 95% rename from tools/pd-simulator/simulator/cases/add_nodes_dynamic.go rename to tools/pd-dev/pd-simulator/simulator/cases/add_nodes_dynamic.go index 59b0b54e1ca..e6b7fd9203d 100644 --- a/tools/pd-simulator/simulator/cases/add_nodes_dynamic.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/add_nodes_dynamic.go @@ -20,8 +20,8 @@ import ( "github.com/docker/go-units" "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/balance_leader.go b/tools/pd-dev/pd-simulator/simulator/cases/balance_leader.go similarity index 93% rename from tools/pd-simulator/simulator/cases/balance_leader.go rename to tools/pd-dev/pd-simulator/simulator/cases/balance_leader.go index bbc7ce97f68..5163cfbe42e 100644 --- a/tools/pd-simulator/simulator/cases/balance_leader.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/balance_leader.go @@ -18,8 +18,8 @@ import ( "github.com/docker/go-units" "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/balance_region.go b/tools/pd-dev/pd-simulator/simulator/cases/balance_region.go similarity index 89% rename from tools/pd-simulator/simulator/cases/balance_region.go rename to tools/pd-dev/pd-simulator/simulator/cases/balance_region.go index 3b0c46f1670..62cd507ee45 100644 --- a/tools/pd-simulator/simulator/cases/balance_region.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/balance_region.go @@ -19,16 +19,17 @@ import ( "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/config" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) func newRedundantBalanceRegion() *Case { var simCase Case - storeNum := simutil.CaseConfigure.StoreNum - regionNum := simutil.CaseConfigure.RegionNum + storeNum := config.CaseConfigure.StoreNum + regionNum := config.CaseConfigure.RegionNum if storeNum == 0 || regionNum == 0 { storeNum, regionNum = 6, 4000 } diff --git a/tools/pd-simulator/simulator/cases/cases.go b/tools/pd-dev/pd-simulator/simulator/cases/cases.go similarity index 91% rename from tools/pd-simulator/simulator/cases/cases.go rename to tools/pd-dev/pd-simulator/simulator/cases/cases.go index 0a8967a8d86..5d79ff685ed 100644 --- a/tools/pd-simulator/simulator/cases/cases.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/cases.go @@ -19,8 +19,9 @@ import ( "github.com/tikv/pd/pkg/core" "github.com/tikv/pd/pkg/schedule/placement" "github.com/tikv/pd/pkg/utils/typeutil" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/config" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" ) // Store is used to simulate tikv. @@ -85,7 +86,7 @@ func (a *idAllocator) GetID() uint64 { // IDAllocator is used to alloc unique ID. var IDAllocator idAllocator -// CaseMap is a mapping of the cases to the their corresponding initialize functions. +// CaseMap is a mapping of the cases to the corresponding to initialize functions. var CaseMap = map[string]func() *Case{ "balance-leader": newBalanceLeader, "redundant-balance-region": newRedundantBalanceRegion, @@ -124,7 +125,7 @@ func isUniform(count, meanCount int, threshold float64) bool { } func getStoreNum() int { - storeNum := simutil.CaseConfigure.StoreNum + storeNum := config.CaseConfigure.StoreNum if storeNum < 3 { simutil.Logger.Fatal("store num should be larger than or equal to 3") } @@ -132,7 +133,7 @@ func getStoreNum() int { } func getRegionNum() int { - regionNum := simutil.CaseConfigure.RegionNum + regionNum := config.CaseConfigure.RegionNum if regionNum <= 0 { simutil.Logger.Fatal("region num should be larger than 0") } diff --git a/tools/pd-simulator/simulator/cases/delete_nodes.go b/tools/pd-dev/pd-simulator/simulator/cases/delete_nodes.go similarity index 95% rename from tools/pd-simulator/simulator/cases/delete_nodes.go rename to tools/pd-dev/pd-simulator/simulator/cases/delete_nodes.go index 4ba8e5064a4..8cef9bd80c2 100644 --- a/tools/pd-simulator/simulator/cases/delete_nodes.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/delete_nodes.go @@ -20,8 +20,8 @@ import ( "github.com/docker/go-units" "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/diagnose_label_isolation.go b/tools/pd-dev/pd-simulator/simulator/cases/diagnose_label_isolation.go similarity index 98% rename from tools/pd-simulator/simulator/cases/diagnose_label_isolation.go rename to tools/pd-dev/pd-simulator/simulator/cases/diagnose_label_isolation.go index 7fa50e56197..8b29d745ee4 100644 --- a/tools/pd-simulator/simulator/cases/diagnose_label_isolation.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/diagnose_label_isolation.go @@ -21,8 +21,8 @@ import ( "github.com/docker/go-units" "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/diagnose_rule.go b/tools/pd-dev/pd-simulator/simulator/cases/diagnose_rule.go similarity index 98% rename from tools/pd-simulator/simulator/cases/diagnose_rule.go rename to tools/pd-dev/pd-simulator/simulator/cases/diagnose_rule.go index 15c5942d810..c1a1c40f370 100644 --- a/tools/pd-simulator/simulator/cases/diagnose_rule.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/diagnose_rule.go @@ -21,8 +21,8 @@ import ( "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" "github.com/tikv/pd/pkg/schedule/placement" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/event_inner.go b/tools/pd-dev/pd-simulator/simulator/cases/event_inner.go similarity index 100% rename from tools/pd-simulator/simulator/cases/event_inner.go rename to tools/pd-dev/pd-simulator/simulator/cases/event_inner.go diff --git a/tools/pd-simulator/simulator/cases/hot_read.go b/tools/pd-dev/pd-simulator/simulator/cases/hot_read.go similarity index 95% rename from tools/pd-simulator/simulator/cases/hot_read.go rename to tools/pd-dev/pd-simulator/simulator/cases/hot_read.go index d4ec6831d95..d23c896e8f4 100644 --- a/tools/pd-simulator/simulator/cases/hot_read.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/hot_read.go @@ -20,8 +20,8 @@ import ( "github.com/docker/go-units" "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/hot_write.go b/tools/pd-dev/pd-simulator/simulator/cases/hot_write.go similarity index 95% rename from tools/pd-simulator/simulator/cases/hot_write.go rename to tools/pd-dev/pd-simulator/simulator/cases/hot_write.go index 8428afa75b5..33ce494b4c2 100644 --- a/tools/pd-simulator/simulator/cases/hot_write.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/hot_write.go @@ -20,8 +20,8 @@ import ( "github.com/docker/go-units" "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/import_data.go b/tools/pd-dev/pd-simulator/simulator/cases/import_data.go similarity index 97% rename from tools/pd-simulator/simulator/cases/import_data.go rename to tools/pd-dev/pd-simulator/simulator/cases/import_data.go index 6cf3b79a736..f265cb77ca5 100644 --- a/tools/pd-simulator/simulator/cases/import_data.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/import_data.go @@ -26,8 +26,8 @@ import ( "github.com/pingcap/log" "github.com/tikv/pd/pkg/codec" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/makeup_down_replica.go b/tools/pd-dev/pd-simulator/simulator/cases/makeup_down_replica.go similarity index 95% rename from tools/pd-simulator/simulator/cases/makeup_down_replica.go rename to tools/pd-dev/pd-simulator/simulator/cases/makeup_down_replica.go index 86c9b4cac1d..6e66fe4072f 100644 --- a/tools/pd-simulator/simulator/cases/makeup_down_replica.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/makeup_down_replica.go @@ -18,8 +18,8 @@ import ( "github.com/docker/go-units" "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/region_merge.go b/tools/pd-dev/pd-simulator/simulator/cases/region_merge.go similarity index 94% rename from tools/pd-simulator/simulator/cases/region_merge.go rename to tools/pd-dev/pd-simulator/simulator/cases/region_merge.go index 3d5d57f804f..e30df0a190e 100644 --- a/tools/pd-simulator/simulator/cases/region_merge.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/region_merge.go @@ -20,8 +20,8 @@ import ( "github.com/docker/go-units" "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/cases/region_split.go b/tools/pd-dev/pd-simulator/simulator/cases/region_split.go similarity index 93% rename from tools/pd-simulator/simulator/cases/region_split.go rename to tools/pd-dev/pd-simulator/simulator/cases/region_split.go index b85cd319494..8c9e9c25670 100644 --- a/tools/pd-simulator/simulator/cases/region_split.go +++ b/tools/pd-dev/pd-simulator/simulator/cases/region_split.go @@ -18,8 +18,8 @@ import ( "github.com/docker/go-units" "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/client.go b/tools/pd-dev/pd-simulator/simulator/client.go similarity index 97% rename from tools/pd-simulator/simulator/client.go rename to tools/pd-dev/pd-simulator/simulator/client.go index 808c991e97f..c960b8254b1 100644 --- a/tools/pd-simulator/simulator/client.go +++ b/tools/pd-dev/pd-simulator/simulator/client.go @@ -30,7 +30,8 @@ import ( "github.com/tikv/pd/pkg/core" "github.com/tikv/pd/pkg/schedule/placement" "github.com/tikv/pd/pkg/utils/typeutil" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/config" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" @@ -45,7 +46,7 @@ type Client interface { PutStore(ctx context.Context, store *metapb.Store) error StoreHeartbeat(ctx context.Context, stats *pdpb.StoreStats) error RegionHeartbeat(ctx context.Context, region *core.RegionInfo) error - PutPDConfig(*PDConfig) error + PutPDConfig(*config.PDConfig) error Close() } @@ -316,7 +317,7 @@ func (c *client) PutStore(ctx context.Context, store *metapb.Store) error { return nil } -func (c *client) PutPDConfig(config *PDConfig) error { +func (c *client) PutPDConfig(config *config.PDConfig) error { if len(config.PlacementRules) > 0 { path := fmt.Sprintf("%s/%s/config/rules/batch", c.url, httpPrefix) ruleOps := make([]*placement.RuleOp, 0) diff --git a/tools/pd-simulator/simulator/simutil/case_config.go b/tools/pd-dev/pd-simulator/simulator/config/case_config.go similarity index 93% rename from tools/pd-simulator/simulator/simutil/case_config.go rename to tools/pd-dev/pd-simulator/simulator/config/case_config.go index a34035c15aa..98cdf90b19e 100644 --- a/tools/pd-simulator/simulator/simutil/case_config.go +++ b/tools/pd-dev/pd-simulator/simulator/config/case_config.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package simutil +package config // CaseConfig is to save flags type CaseConfig struct { @@ -21,7 +21,7 @@ type CaseConfig struct { EnableTransferRegionCounter bool } -// CaseConfigure is an global instance for CaseConfig +// CaseConfigure is a global instance for CaseConfig var CaseConfigure *CaseConfig // InitCaseConfig is to init caseConfigure diff --git a/tools/pd-simulator/simulator/config.go b/tools/pd-dev/pd-simulator/simulator/config/config.go similarity index 98% rename from tools/pd-simulator/simulator/config.go rename to tools/pd-dev/pd-simulator/simulator/config/config.go index 4f197fb83c2..ddd9d9b20e8 100644 --- a/tools/pd-simulator/simulator/config.go +++ b/tools/pd-dev/pd-simulator/simulator/config/config.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package simulator +package config import ( "fmt" @@ -118,7 +118,7 @@ func (sc *SimConfig) Adjust(meta *toml.MetaData) error { return sc.ServerConfig.Adjust(meta, false) } -func (sc *SimConfig) speed() uint64 { +func (sc *SimConfig) Speed() uint64 { return uint64(time.Second / sc.SimTickInterval.Duration) } diff --git a/tools/pd-simulator/simulator/conn.go b/tools/pd-dev/pd-simulator/simulator/conn.go similarity index 89% rename from tools/pd-simulator/simulator/conn.go rename to tools/pd-dev/pd-simulator/simulator/conn.go index 588fec246d4..6ecab704693 100644 --- a/tools/pd-simulator/simulator/conn.go +++ b/tools/pd-dev/pd-simulator/simulator/conn.go @@ -16,7 +16,8 @@ package simulator import ( "github.com/pingcap/kvproto/pkg/metapb" - "github.com/tikv/pd/tools/pd-simulator/simulator/cases" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/cases" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/config" ) // Connection records the information of connection among nodes. @@ -26,7 +27,7 @@ type Connection struct { } // NewConnection creates nodes according to the configuration and returns the connection among nodes. -func NewConnection(simCase *cases.Case, pdAddr string, storeConfig *SimConfig) (*Connection, error) { +func NewConnection(simCase *cases.Case, pdAddr string, storeConfig *config.SimConfig) (*Connection, error) { conn := &Connection{ pdAddr: pdAddr, Nodes: make(map[uint64]*Node), diff --git a/tools/pd-simulator/simulator/drive.go b/tools/pd-dev/pd-simulator/simulator/drive.go similarity index 92% rename from tools/pd-simulator/simulator/drive.go rename to tools/pd-dev/pd-simulator/simulator/drive.go index c7f64324c19..9bb3dcee7f2 100644 --- a/tools/pd-simulator/simulator/drive.go +++ b/tools/pd-dev/pd-simulator/simulator/drive.go @@ -25,9 +25,10 @@ import ( "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" "github.com/tikv/pd/pkg/utils/typeutil" - "github.com/tikv/pd/tools/pd-simulator/simulator/cases" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/cases" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/config" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.etcd.io/etcd/clientv3" "go.uber.org/zap" ) @@ -42,17 +43,17 @@ type Driver struct { eventRunner *EventRunner raftEngine *RaftEngine conn *Connection - simConfig *SimConfig - pdConfig *PDConfig + simConfig *config.SimConfig + pdConfig *config.PDConfig } // NewDriver returns a driver. -func NewDriver(pdAddr string, caseName string, simConfig *SimConfig) (*Driver, error) { +func NewDriver(pdAddr string, caseName string, simConfig *config.SimConfig) (*Driver, error) { simCase := cases.NewCase(caseName) if simCase == nil { return nil, errors.Errorf("failed to create case %s", caseName) } - pdConfig := &PDConfig{} + pdConfig := &config.PDConfig{} pdConfig.PlacementRules = simCase.Rules pdConfig.LocationLabels = simCase.Labels return &Driver{ diff --git a/tools/pd-simulator/simulator/event.go b/tools/pd-dev/pd-simulator/simulator/event.go similarity index 97% rename from tools/pd-simulator/simulator/event.go rename to tools/pd-dev/pd-simulator/simulator/event.go index 04ad10a0db8..578840cb9f2 100644 --- a/tools/pd-simulator/simulator/event.go +++ b/tools/pd-dev/pd-simulator/simulator/event.go @@ -18,8 +18,8 @@ import ( "github.com/pingcap/kvproto/pkg/metapb" "github.com/pingcap/kvproto/pkg/pdpb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-simulator/simulator/cases" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/cases" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) diff --git a/tools/pd-simulator/simulator/info/stats.go b/tools/pd-dev/pd-simulator/simulator/info/stats.go similarity index 100% rename from tools/pd-simulator/simulator/info/stats.go rename to tools/pd-dev/pd-simulator/simulator/info/stats.go diff --git a/tools/pd-simulator/simulator/metrics.go b/tools/pd-dev/pd-simulator/simulator/metrics.go similarity index 100% rename from tools/pd-simulator/simulator/metrics.go rename to tools/pd-dev/pd-simulator/simulator/metrics.go diff --git a/tools/pd-simulator/simulator/node.go b/tools/pd-dev/pd-simulator/simulator/node.go similarity index 95% rename from tools/pd-simulator/simulator/node.go rename to tools/pd-dev/pd-simulator/simulator/node.go index 68a10a8638e..76f820325f0 100644 --- a/tools/pd-simulator/simulator/node.go +++ b/tools/pd-dev/pd-simulator/simulator/node.go @@ -26,9 +26,10 @@ import ( "github.com/pingcap/kvproto/pkg/pdpb" "github.com/tikv/pd/pkg/ratelimit" "github.com/tikv/pd/pkg/utils/syncutil" - "github.com/tikv/pd/tools/pd-simulator/simulator/cases" - "github.com/tikv/pd/tools/pd-simulator/simulator/info" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/cases" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/config" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/info" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) @@ -57,7 +58,7 @@ type Node struct { } // NewNode returns a Node. -func NewNode(s *cases.Store, pdAddr string, config *SimConfig) (*Node, error) { +func NewNode(s *cases.Store, pdAddr string, config *config.SimConfig) (*Node, error) { ctx, cancel := context.WithCancel(context.Background()) store := &metapb.Store{ Id: s.ID, @@ -93,7 +94,7 @@ func NewNode(s *cases.Store, pdAddr string, config *SimConfig) (*Node, error) { cancel() return nil, err } - ratio := config.speed() + ratio := config.Speed() speed := config.StoreIOMBPerSecond * units.MiB * int64(ratio) return &Node{ Store: store, diff --git a/tools/pd-simulator/simulator/raft.go b/tools/pd-dev/pd-simulator/simulator/raft.go similarity index 96% rename from tools/pd-simulator/simulator/raft.go rename to tools/pd-dev/pd-simulator/simulator/raft.go index fccf75781d3..4fb3b022509 100644 --- a/tools/pd-simulator/simulator/raft.go +++ b/tools/pd-dev/pd-simulator/simulator/raft.go @@ -21,8 +21,9 @@ import ( "github.com/pingcap/kvproto/pkg/metapb" "github.com/tikv/pd/pkg/core" "github.com/tikv/pd/pkg/utils/syncutil" - "github.com/tikv/pd/tools/pd-simulator/simulator/cases" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/cases" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/config" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) @@ -34,12 +35,12 @@ type RaftEngine struct { regionChange map[uint64][]uint64 regionSplitSize int64 regionSplitKeys int64 - storeConfig *SimConfig + storeConfig *config.SimConfig useTiDBEncodedKey bool } // NewRaftEngine creates the initialized raft with the configuration. -func NewRaftEngine(conf *cases.Case, conn *Connection, storeConfig *SimConfig) *RaftEngine { +func NewRaftEngine(conf *cases.Case, conn *Connection, storeConfig *config.SimConfig) *RaftEngine { r := &RaftEngine{ regionsInfo: core.NewRegionsInfo(), conn: conn, diff --git a/tools/pd-simulator/simulator/simutil/key.go b/tools/pd-dev/pd-simulator/simulator/simutil/key.go similarity index 100% rename from tools/pd-simulator/simulator/simutil/key.go rename to tools/pd-dev/pd-simulator/simulator/simutil/key.go diff --git a/tools/pd-simulator/simulator/simutil/key_test.go b/tools/pd-dev/pd-simulator/simulator/simutil/key_test.go similarity index 100% rename from tools/pd-simulator/simulator/simutil/key_test.go rename to tools/pd-dev/pd-simulator/simulator/simutil/key_test.go diff --git a/tools/pd-simulator/simulator/simutil/logger.go b/tools/pd-dev/pd-simulator/simulator/simutil/logger.go similarity index 100% rename from tools/pd-simulator/simulator/simutil/logger.go rename to tools/pd-dev/pd-simulator/simulator/simutil/logger.go diff --git a/tools/pd-simulator/simulator/task.go b/tools/pd-dev/pd-simulator/simulator/task.go similarity index 99% rename from tools/pd-simulator/simulator/task.go rename to tools/pd-dev/pd-simulator/simulator/task.go index a19854b53ba..e8a6c9a0774 100644 --- a/tools/pd-simulator/simulator/task.go +++ b/tools/pd-dev/pd-simulator/simulator/task.go @@ -25,8 +25,8 @@ import ( "github.com/pingcap/kvproto/pkg/metapb" "github.com/pingcap/kvproto/pkg/pdpb" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/tools/pd-analysis/analysis" - "github.com/tikv/pd/tools/pd-simulator/simulator/simutil" + "github.com/tikv/pd/tools/pd-dev/pd-analysis/analysis" + "github.com/tikv/pd/tools/pd-dev/pd-simulator/simulator/simutil" "go.uber.org/zap" ) @@ -415,7 +415,7 @@ func (a *addPeer) tick(engine *RaftEngine, region *core.RegionInfo) (newRegion * pendingPeers := append(region.GetPendingPeers(), a.peer) return region.Clone(core.WithAddPeer(a.peer), core.WithIncConfVer(), core.WithPendingPeers(pendingPeers)), false } - speed := engine.storeConfig.speed() + speed := engine.storeConfig.Speed() // Step 2: Process Snapshot if !processSnapshot(sendNode, a.sendingStat, speed) { return nil, false diff --git a/tools/pd-tso-bench/README.md b/tools/pd-dev/pd-tso-bench/README.md similarity index 100% rename from tools/pd-tso-bench/README.md rename to tools/pd-dev/pd-tso-bench/README.md diff --git a/tools/pd-dev/pd-tso-bench/config.go b/tools/pd-dev/pd-tso-bench/config.go new file mode 100644 index 00000000000..1c1f747046d --- /dev/null +++ b/tools/pd-dev/pd-tso-bench/config.go @@ -0,0 +1,70 @@ +// Copyright 2024 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 tsobench + +import ( + "time" + + flag "github.com/spf13/pflag" + "github.com/tikv/pd/tools/pd-dev/util" +) + +// Config is the heartbeat-bench configuration. +type config struct { + *util.GeneralConfig + Client int `toml:"client" json:"client"` + Duration time.Duration `toml:"duration" json:"duration"` + Count int `toml:"count" json:"count"` + Concurrency int `toml:"concurrency" json:"concurrency"` + DcLocation string `toml:"dc-location" json:"dc-location"` + Verbose bool `toml:"verbose" json:"verbose"` + Interval time.Duration `toml:"interval" json:"interval"` + + MaxBatchWaitInterval time.Duration `toml:"max-batch-wait-interval" json:"max-batch-wait-interval"` + EnableTSOFollowerProxy bool `toml:"enable-tso-follower-proxy" json:"enable-tso-follower-proxy"` + EnableFaultInjection bool `toml:"enable-fault-injection" json:"enable-fault-injection"` + FaultInjectionRate float64 `toml:"fault-injection-rate" json:"fault-injection-rate"` + MaxTSOSendIntervalMilliseconds int `toml:"max-tso-send-interval-ms" json:"max-tso-send-interval-ms"` + KeyspaceID uint `toml:"keyspace-id" json:"keyspace-id"` + KeyspaceName string `toml:"keyspace-name" json:"keyspace-name"` +} + +// NewConfig return a set of settings. +func newConfig() *config { + cfg := &config{ + GeneralConfig: &util.GeneralConfig{}, + } + cfg.FlagSet = flag.NewFlagSet("api-bench", flag.ContinueOnError) + fs := cfg.FlagSet + cfg.GeneralConfig = util.NewGeneralConfig(fs) + fs.ParseErrorsWhitelist.UnknownFlags = true + + fs.Int("client", 1, "the number of pd clients involved in each benchmark") + fs.Int("concurrency", 1000, "concurrency") + fs.Int("count", 1, "the count number that the test will run") + fs.Duration("duration", 60*time.Second, "how many seconds the test will last") + fs.String("dc", "global", "which dc-location this bench will request") + fs.Bool("v", false, "output statistics info every interval and output metrics info at the end") + fs.Duration("interval", time.Second, "interval to output the statistics") + fs.Duration("batch-interval", 0, "the max batch wait interval") + fs.Bool("enable-tso-follower-proxy", false, "whether enable the TSO Follower Proxy") + fs.Bool("enable-fault-injection", false, "whether enable fault injection") + fs.Float64("fault-injection-rate", 0.01, "the failure rate [0.0001, 1]. 0.01 means 1% failure rate") + fs.Int("max-send-interval-ms", 0, "max tso send interval in milliseconds, 60s by default") + fs.Uint("keyspace-id", 0, "the id of the keyspace to access") + fs.String("keyspace-name", "", "the name of the keyspace to access") + + return cfg +} diff --git a/tools/pd-dev/pd-tso-bench/main.go b/tools/pd-dev/pd-tso-bench/main.go new file mode 100644 index 00000000000..9b8803b1601 --- /dev/null +++ b/tools/pd-dev/pd-tso-bench/main.go @@ -0,0 +1,247 @@ +// Copyright 2017 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 tsobench + +import ( + "context" + "fmt" + "io" + "math/rand" + "net/http" + "net/http/httptest" + "os" + "sync" + "time" + + "github.com/influxdata/tdigest" + "github.com/pingcap/errors" + "github.com/pingcap/log" + "github.com/prometheus/client_golang/prometheus/promhttp" + "github.com/spf13/pflag" + pd "github.com/tikv/pd/client" + "github.com/tikv/pd/pkg/utils/logutil" + "go.uber.org/zap" + "google.golang.org/grpc" + "google.golang.org/grpc/keepalive" +) + +const ( + keepaliveTime = 10 * time.Second + keepaliveTimeout = 3 * time.Second +) + +var ( + wg sync.WaitGroup + promServer *httptest.Server +) + +func collectMetrics(server *httptest.Server) string { + time.Sleep(1100 * time.Millisecond) + res, _ := http.Get(server.URL) + body, _ := io.ReadAll(res.Body) + res.Body.Close() + return string(body) +} + +func Run(ctx context.Context) { + defer logutil.LogPanic() + + cfg := newConfig() + err := cfg.Parse(os.Args[1:]) + switch errors.Cause(err) { + case nil: + case pflag.ErrHelp: + os.Exit(0) + default: + log.Fatal("parse cmd flags error", zap.Error(err)) + } + + for i := 0; i < cfg.Count; i++ { + fmt.Printf("\nStart benchmark #%d, duration: %+vs\n", i, cfg.Duration.Seconds()) + bench(ctx, cfg) + } +} + +func bench(mainCtx context.Context, cfg *config) { + promServer = httptest.NewServer(promhttp.Handler()) + + // Initialize all clients + fmt.Printf("Create %d client(s) for benchmark\n", cfg.Client) + pdClients := make([]pd.Client, cfg.Client) + for idx := range pdClients { + pdCli, err := createPDClient(mainCtx, cfg) + if err != nil { + log.Fatal(fmt.Sprintf("create pd client #%d failed: %v", idx, err)) + } + pdClients[idx] = pdCli + } + + ctx, cancel := context.WithCancel(mainCtx) + // To avoid the first time high latency. + for idx, pdCli := range pdClients { + _, _, err := pdCli.GetLocalTS(ctx, cfg.DcLocation) + if err != nil { + log.Fatal("get first time tso failed", zap.Int("client-number", idx), zap.Error(err)) + } + } + + durCh := make(chan time.Duration, 2*(cfg.Concurrency)*(cfg.Client)) + + if cfg.EnableFaultInjection { + fmt.Printf("Enable fault injection, failure rate: %f\n", cfg.FaultInjectionRate) + wg.Add(cfg.Client) + for i := 0; i < cfg.Client; i++ { + go reqWorker(ctx, cfg, pdClients, i, durCh) + } + } else { + wg.Add(cfg.Concurrency * cfg.Client) + for i := 0; i < cfg.Client; i++ { + for j := 0; j < cfg.Concurrency; j++ { + go reqWorker(ctx, cfg, pdClients, i, durCh) + } + } + } + + wg.Add(1) + go showStats(ctx, cfg.Interval, cfg.Verbose, durCh) + + timer := time.NewTimer(cfg.Duration) + defer timer.Stop() + + select { + case <-ctx.Done(): + case <-timer.C: + } + cancel() + + wg.Wait() + for _, pdCli := range pdClients { + pdCli.Close() + } +} + +var latencyTDigest = tdigest.New() + +func reqWorker(ctx context.Context, cfg *config, + pdClients []pd.Client, clientIdx int, durCh chan time.Duration) { + defer wg.Done() + + reqCtx, cancel := context.WithCancel(ctx) + defer cancel() + var ( + err error + maxRetryTime = 120 + sleepIntervalOnFailure = 1000 * time.Millisecond + totalSleepBeforeGetTS time.Duration + ) + pdCli := pdClients[clientIdx] + + for { + if pdCli == nil || (cfg.EnableFaultInjection && shouldInjectFault(cfg)) { + if pdCli != nil { + pdCli.Close() + } + pdCli, err = createPDClient(ctx, cfg) + if err != nil { + log.Error(fmt.Sprintf("re-create pd client #%d failed: %v", clientIdx, err)) + select { + case <-reqCtx.Done(): + case <-time.After(100 * time.Millisecond): + } + continue + } + pdClients[clientIdx] = pdCli + } + + totalSleepBeforeGetTS = 0 + start := time.Now() + + i := 0 + for ; i < maxRetryTime; i++ { + var ticker *time.Ticker + if cfg.MaxTSOSendIntervalMilliseconds > 0 { + sleepBeforeGetTS := time.Duration(rand.Intn(cfg.MaxTSOSendIntervalMilliseconds)) * time.Millisecond + ticker = time.NewTicker(sleepBeforeGetTS) + select { + case <-reqCtx.Done(): + case <-ticker.C: + totalSleepBeforeGetTS += sleepBeforeGetTS + } + } + _, _, err = pdCli.GetLocalTS(reqCtx, cfg.DcLocation) + if errors.Cause(err) == context.Canceled { + ticker.Stop() + return + } + if err == nil { + ticker.Stop() + break + } + log.Error(fmt.Sprintf("%v", err)) + time.Sleep(sleepIntervalOnFailure) + } + if err != nil { + log.Fatal(fmt.Sprintf("%v", err)) + } + dur := time.Since(start) - time.Duration(i)*sleepIntervalOnFailure - totalSleepBeforeGetTS + + select { + case <-reqCtx.Done(): + return + case durCh <- dur: + } + } +} + +func createPDClient(ctx context.Context, cfg *config) (pd.Client, error) { + var ( + pdCli pd.Client + err error + ) + + opts := make([]pd.ClientOption, 0) + opts = append(opts, pd.WithGRPCDialOptions( + grpc.WithKeepaliveParams(keepalive.ClientParameters{ + Time: keepaliveTime, + Timeout: keepaliveTimeout, + }), + )) + + if len(cfg.KeyspaceName) > 0 { + apiCtx := pd.NewAPIContextV2(cfg.KeyspaceName) + pdCli, err = pd.NewClientWithAPIContext(ctx, apiCtx, []string{cfg.PDAddrs}, pd.SecurityOption{ + CAPath: cfg.CaPath, + CertPath: cfg.CertPath, + KeyPath: cfg.KeyPath, + }, opts...) + } else { + pdCli, err = pd.NewClientWithKeyspace(ctx, uint32(cfg.KeyspaceID), []string{cfg.PDAddrs}, pd.SecurityOption{ + CAPath: cfg.CaPath, + CertPath: cfg.CertPath, + KeyPath: cfg.KeyPath, + }, opts...) + } + if err != nil { + return nil, err + } + + pdCli.UpdateOption(pd.MaxTSOBatchWaitInterval, cfg.MaxBatchWaitInterval) + pdCli.UpdateOption(pd.EnableTSOFollowerProxy, cfg.EnableTSOFollowerProxy) + return pdCli, err +} + +func shouldInjectFault(cfg *config) bool { + return rand.Intn(10000) < int(cfg.FaultInjectionRate*10000) +} diff --git a/tools/pd-dev/pd-tso-bench/stats.go b/tools/pd-dev/pd-tso-bench/stats.go new file mode 100644 index 00000000000..cb46377a19a --- /dev/null +++ b/tools/pd-dev/pd-tso-bench/stats.go @@ -0,0 +1,211 @@ +// Copyright 2024 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 tsobench + +import ( + "context" + "fmt" + "time" +) + +const ( + twoDur = time.Millisecond * 2 + fiveDur = time.Millisecond * 5 + tenDur = time.Millisecond * 10 + thirtyDur = time.Millisecond * 30 + fiftyDur = time.Millisecond * 50 + oneHundredDur = time.Millisecond * 100 + twoHundredDur = time.Millisecond * 200 + fourHundredDur = time.Millisecond * 400 + eightHundredDur = time.Millisecond * 800 + oneThousandDur = time.Millisecond * 1000 +) + +type stats struct { + maxDur time.Duration + minDur time.Duration + totalDur time.Duration + count int + submilliCnt int + milliCnt int + twoMilliCnt int + fiveMilliCnt int + tenMSCnt int + thirtyCnt int + fiftyCnt int + oneHundredCnt int + twoHundredCnt int + fourHundredCnt int + eightHundredCnt int + oneThousandCnt int +} + +func newStats() *stats { + return &stats{ + minDur: time.Hour, + maxDur: 0, + } +} + +func (s *stats) update(dur time.Duration) { + s.count++ + s.totalDur += dur + latencyTDigest.Add(float64(dur.Nanoseconds())/1e6, 1) + + if dur > s.maxDur { + s.maxDur = dur + } + if dur < s.minDur { + s.minDur = dur + } + + if dur > oneThousandDur { + s.oneThousandCnt++ + return + } + + if dur > eightHundredDur { + s.eightHundredCnt++ + return + } + + if dur > fourHundredDur { + s.fourHundredCnt++ + return + } + + if dur > twoHundredDur { + s.twoHundredCnt++ + return + } + + if dur > oneHundredDur { + s.oneHundredCnt++ + return + } + + if dur > fiftyDur { + s.fiftyCnt++ + return + } + + if dur > thirtyDur { + s.thirtyCnt++ + return + } + + if dur > tenDur { + s.tenMSCnt++ + return + } + + if dur > fiveDur { + s.fiveMilliCnt++ + return + } + + if dur > twoDur { + s.twoMilliCnt++ + return + } + + if dur > time.Millisecond { + s.milliCnt++ + return + } + + s.submilliCnt++ +} + +func (s *stats) merge(other *stats) { + if s.maxDur < other.maxDur { + s.maxDur = other.maxDur + } + if s.minDur > other.minDur { + s.minDur = other.minDur + } + + s.count += other.count + s.totalDur += other.totalDur + s.submilliCnt += other.submilliCnt + s.milliCnt += other.milliCnt + s.twoMilliCnt += other.twoMilliCnt + s.fiveMilliCnt += other.fiveMilliCnt + s.tenMSCnt += other.tenMSCnt + s.thirtyCnt += other.thirtyCnt + s.fiftyCnt += other.fiftyCnt + s.oneHundredCnt += other.oneHundredCnt + s.twoHundredCnt += other.twoHundredCnt + s.fourHundredCnt += other.fourHundredCnt + s.eightHundredCnt += other.eightHundredCnt + s.oneThousandCnt += other.oneThousandCnt +} + +func (s *stats) Counter() string { + return fmt.Sprintf( + "count: %d, max: %.4fms, min: %.4fms, avg: %.4fms\n<1ms: %d, >1ms: %d, >2ms: %d, >5ms: %d, >10ms: %d, >30ms: %d, >50ms: %d, >100ms: %d, >200ms: %d, >400ms: %d, >800ms: %d, >1s: %d", + s.count, float64(s.maxDur.Nanoseconds())/float64(time.Millisecond), float64(s.minDur.Nanoseconds())/float64(time.Millisecond), float64(s.totalDur.Nanoseconds())/float64(s.count)/float64(time.Millisecond), + s.submilliCnt, s.milliCnt, s.twoMilliCnt, s.fiveMilliCnt, s.tenMSCnt, s.thirtyCnt, s.fiftyCnt, s.oneHundredCnt, s.twoHundredCnt, s.fourHundredCnt, + s.eightHundredCnt, s.oneThousandCnt) +} + +func (s *stats) Percentage() string { + return fmt.Sprintf( + "count: %d, <1ms: %2.2f%%, >1ms: %2.2f%%, >2ms: %2.2f%%, >5ms: %2.2f%%, >10ms: %2.2f%%, >30ms: %2.2f%%, >50ms: %2.2f%%, >100ms: %2.2f%%, >200ms: %2.2f%%, >400ms: %2.2f%%, >800ms: %2.2f%%, >1s: %2.2f%%", s.count, + s.calculate(s.submilliCnt), s.calculate(s.milliCnt), s.calculate(s.twoMilliCnt), s.calculate(s.fiveMilliCnt), s.calculate(s.tenMSCnt), s.calculate(s.thirtyCnt), s.calculate(s.fiftyCnt), + s.calculate(s.oneHundredCnt), s.calculate(s.twoHundredCnt), s.calculate(s.fourHundredCnt), s.calculate(s.eightHundredCnt), s.calculate(s.oneThousandCnt)) +} + +func (s *stats) calculate(count int) float64 { + return float64(count) * 100 / float64(s.count) +} + +func showStats(ctx context.Context, interval time.Duration, verbose bool, durCh chan time.Duration) { + defer wg.Done() + + statCtx, cancel := context.WithCancel(ctx) + defer cancel() + + ticker := time.NewTicker(interval) + defer ticker.Stop() + + s := newStats() + total := newStats() + + fmt.Println() + for { + select { + case <-ticker.C: + // runtime.GC() + if verbose { + fmt.Println(s.Counter()) + } + total.merge(s) + s = newStats() + case d := <-durCh: + s.update(d) + case <-statCtx.Done(): + fmt.Println("\nTotal:") + fmt.Println(total.Counter()) + fmt.Println(total.Percentage()) + // Calculate the percentiles by using the tDigest algorithm. + fmt.Printf("P0.5: %.4fms, P0.8: %.4fms, P0.9: %.4fms, P0.99: %.4fms\n\n", latencyTDigest.Quantile(0.5), latencyTDigest.Quantile(0.8), latencyTDigest.Quantile(0.9), latencyTDigest.Quantile(0.99)) + if verbose { + fmt.Println(collectMetrics(promServer)) + } + return + } + } +} diff --git a/tools/pd-ut/README.md b/tools/pd-dev/pd-ut/README.md similarity index 100% rename from tools/pd-ut/README.md rename to tools/pd-dev/pd-ut/README.md diff --git a/tools/pd-ut/go-compile-without-link.sh b/tools/pd-dev/pd-ut/go-compile-without-link.sh similarity index 100% rename from tools/pd-ut/go-compile-without-link.sh rename to tools/pd-dev/pd-ut/go-compile-without-link.sh diff --git a/tools/pd-ut/ut.go b/tools/pd-dev/pd-ut/ut.go similarity index 97% rename from tools/pd-ut/ut.go rename to tools/pd-dev/pd-ut/ut.go index 7fc96ee11cf..e43fd0530c5 100644 --- a/tools/pd-ut/ut.go +++ b/tools/pd-dev/pd-ut/ut.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package ut import ( "bytes" @@ -92,7 +92,7 @@ var ( junitFile string ) -func main() { +func Run() { race = handleFlag("--race") junitFile = stripFlag("--junitfile") @@ -108,18 +108,19 @@ func main() { var isSucceed bool // run all tests - if len(os.Args) == 1 { + if len(os.Args) == 3 { isSucceed = cmdRun() } - if len(os.Args) >= 2 { - switch os.Args[1] { + if len(os.Args) >= 4 { + // ./pd-dev --mode ut [args] + switch os.Args[3] { case "list": - isSucceed = cmdList(os.Args[2:]...) + isSucceed = cmdList(os.Args[4:]...) case "build": - isSucceed = cmdBuild(os.Args[2:]...) + isSucceed = cmdBuild(os.Args[4:]...) case "run": - isSucceed = cmdRun(os.Args[2:]...) + isSucceed = cmdRun(os.Args[4:]...) default: isSucceed = usage() } @@ -341,9 +342,9 @@ func cmdRun(args ...string) bool { // The value of the flag is returned. func stripFlag(flag string) string { var res string - tmp := os.Args[:0] + tmp := os.Args[:1] // Iter to the flag - var i int + i := 1 for ; i < len(os.Args); i++ { if os.Args[i] == flag { i++ @@ -366,8 +367,8 @@ func stripFlag(flag string) string { } func handleFlag(f string) (found bool) { - tmp := os.Args[:0] - for i := 0; i < len(os.Args); i++ { + tmp := os.Args[:1] + for i := 1; i < len(os.Args); i++ { if os.Args[i] == f { found = true continue @@ -592,7 +593,7 @@ func skipDIR(pkg string) bool { func generateBuildCache() error { // cd cmd/pd-server && go test -tags=tso_function_test,deadlock -exec-=true -vet=off -toolexec=go-compile-without-link cmd := exec.Command("go", "test", "-exec=true", "-vet", "off", "--tags=tso_function_test,deadlock") - goCompileWithoutLink := fmt.Sprintf("-toolexec=%s/tools/pd-ut/go-compile-without-link.sh", workDir) + goCompileWithoutLink := fmt.Sprintf("-toolexec=%s/tools/pd-dev/pd-ut/go-compile-without-link.sh", workDir) cmd.Args = append(cmd.Args, goCompileWithoutLink) cmd.Dir = fmt.Sprintf("%s/cmd/pd-server", workDir) cmd.Stdout = os.Stdout diff --git a/tools/pd-ut/xprog.go b/tools/pd-dev/pd-ut/xprog.go similarity index 100% rename from tools/pd-ut/xprog.go rename to tools/pd-dev/pd-ut/xprog.go diff --git a/tools/pd-dev/regions-dump/config.go b/tools/pd-dev/regions-dump/config.go new file mode 100644 index 00000000000..10737cc3c8b --- /dev/null +++ b/tools/pd-dev/regions-dump/config.go @@ -0,0 +1,46 @@ +// Copyright 2024 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 regiondump + +import ( + flag "github.com/spf13/pflag" + "github.com/tikv/pd/tools/pd-dev/util" +) + +// Config is the heartbeat-bench configuration. +type Config struct { + *util.GeneralConfig + ClusterID uint64 `toml:"cluster-id" json:"cluster-id"` + FilePath string `toml:"file-path" json:"file-path"` + StartID uint64 `toml:"start-id" json:"start-id"` + EndID uint64 `toml:"end-id" json:"end-id"` +} + +// NewConfig return a set of settings. +func newConfig() *Config { + cfg := &Config{ + GeneralConfig: &util.GeneralConfig{}, + } + cfg.FlagSet = flag.NewFlagSet("stores-dump", flag.ContinueOnError) + fs := cfg.FlagSet + cfg.GeneralConfig = util.NewGeneralConfig(fs) + fs.ParseErrorsWhitelist.UnknownFlags = true + fs.Uint64Var(&cfg.ClusterID, "cluster-id", 0, "please make cluster ID match with TiKV") + fs.StringVar(&cfg.FilePath, "file", "stores.dump", "dump file path and name") + fs.Uint64Var(&cfg.StartID, "start-id", 0, "ID of the start region") + fs.Uint64Var(&cfg.EndID, "end-id", 0, "ID of the last region") + + return cfg +} diff --git a/tools/regions-dump/main.go b/tools/pd-dev/regions-dump/main.go similarity index 56% rename from tools/regions-dump/main.go rename to tools/pd-dev/regions-dump/main.go index dddbcbba854..11cc16bf665 100644 --- a/tools/regions-dump/main.go +++ b/tools/pd-dev/regions-dump/main.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package regiondump import ( "bufio" @@ -27,21 +27,12 @@ import ( "github.com/pingcap/errors" "github.com/pingcap/kvproto/pkg/metapb" + "github.com/pingcap/log" "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/pkg/utils/etcdutil" + "github.com/tikv/pd/tools/pd-dev/util" "go.etcd.io/etcd/clientv3" "go.etcd.io/etcd/pkg/transport" -) - -var ( - clusterID = flag.Uint64("cluster-id", 0, "please make cluster ID match with TiKV") - endpoints = flag.String("endpoints", "http://127.0.0.1:2379", "endpoints urls") - startID = flag.Uint64("start-id", 0, "ID of the start region") - endID = flag.Uint64("end-id", 0, "ID of the last region") - filePath = flag.String("file", "regions.dump", "dump file path and name") - caPath = flag.String("cacert", "", "path of file that contains list of trusted SSL CAs") - certPath = flag.String("cert", "", "path of file that contains X509 certificate in PEM format") - keyPath = flag.String("key", "", "path of file that contains X509 key in PEM format") + "go.uber.org/zap" ) const ( @@ -63,13 +54,22 @@ func checkErr(err error) { } } -func main() { - flag.Parse() - if *endID != 0 && *endID < *startID { +func Run() { + cfg := newConfig() + err := cfg.Parse(os.Args[1:]) + switch errors.Cause(err) { + case nil: + case flag.ErrHelp: + os.Exit(0) + default: + log.Fatal("parse cmd flags error", zap.Error(err)) + } + + if cfg.EndID != 0 && cfg.EndID < cfg.StartID { checkErr(errors.New("The end id should great or equal than start id")) } - rootPath = path.Join(pdRootPath, strconv.FormatUint(*clusterID, 10)) - f, err := os.Create(*filePath) + rootPath = path.Join(pdRootPath, strconv.FormatUint(cfg.ClusterID, 10)) + f, err := os.Create(cfg.FilePath) checkErr(err) defer func() { if err := f.Close(); err != nil { @@ -77,12 +77,12 @@ func main() { } }() - urls := strings.Split(*endpoints, ",") + urls := strings.Split(cfg.PDAddrs, ",") tlsInfo := transport.TLSInfo{ - CertFile: *certPath, - KeyFile: *keyPath, - TrustedCAFile: *caPath, + CertFile: cfg.CertPath, + KeyFile: cfg.KeyPath, + TrustedCAFile: cfg.CaPath, } tlsConfig, err := tlsInfo.ClientConfig() checkErr(err) @@ -94,7 +94,7 @@ func main() { }) checkErr(err) - err = loadRegions(client, f) + err = loadRegions(client, cfg, f) checkErr(err) fmt.Println("successful!") } @@ -103,11 +103,11 @@ func regionPath(regionID uint64) string { return path.Join("raft", "r", fmt.Sprintf("%020d", regionID)) } -func loadRegions(client *clientv3.Client, f *os.File) error { - nextID := *startID +func loadRegions(client *clientv3.Client, cfg *Config, f *os.File) error { + nextID := cfg.StartID endKey := regionPath(math.MaxUint64) - if *endID != 0 { - endKey = regionPath(*endID) + if cfg.EndID != 0 { + endKey = regionPath(cfg.EndID) } w := bufio.NewWriter(f) defer w.Flush() @@ -117,7 +117,7 @@ func loadRegions(client *clientv3.Client, f *os.File) error { rangeLimit := maxKVRangeLimit for { startKey := regionPath(nextID) - _, res, err := loadRange(client, startKey, endKey, rangeLimit) + _, res, err := util.LoadRange(client, rootPath, startKey, endKey, rangeLimit) if err != nil { if rangeLimit /= 2; rangeLimit >= minKVRangeLimit { continue @@ -139,22 +139,3 @@ func loadRegions(client *clientv3.Client, f *os.File) error { } } } - -func loadRange(client *clientv3.Client, key, endKey string, limit int) ([]string, []string, error) { - key = path.Join(rootPath, key) - endKey = path.Join(rootPath, endKey) - - withRange := clientv3.WithRange(endKey) - withLimit := clientv3.WithLimit(int64(limit)) - resp, err := etcdutil.EtcdKVGet(client, key, withRange, withLimit) - if err != nil { - return nil, nil, err - } - keys := make([]string, 0, len(resp.Kvs)) - values := make([]string, 0, len(resp.Kvs)) - for _, item := range resp.Kvs { - keys = append(keys, strings.TrimPrefix(strings.TrimPrefix(string(item.Key), rootPath), "/")) - values = append(values, string(item.Value)) - } - return keys, values, nil -} diff --git a/tools/pd-dev/stores-dump/config.go b/tools/pd-dev/stores-dump/config.go new file mode 100644 index 00000000000..7d344b8cf10 --- /dev/null +++ b/tools/pd-dev/stores-dump/config.go @@ -0,0 +1,42 @@ +// Copyright 2024 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 storesdump + +import ( + flag "github.com/spf13/pflag" + "github.com/tikv/pd/tools/pd-dev/util" +) + +// Config is the heartbeat-bench configuration. +type Config struct { + *util.GeneralConfig + ClusterID uint64 `toml:"cluster-id" json:"cluster-id"` + FilePath string `toml:"file" json:"file"` +} + +// NewConfig return a set of settings. +func newConfig() *Config { + cfg := &Config{ + GeneralConfig: &util.GeneralConfig{}, + } + cfg.FlagSet = flag.NewFlagSet("stores-dump", flag.ContinueOnError) + fs := cfg.FlagSet + cfg.GeneralConfig = util.NewGeneralConfig(fs) + fs.ParseErrorsWhitelist.UnknownFlags = true + fs.Uint64Var(&cfg.ClusterID, "cluster-id", 0, "please make cluster ID match with TiKV") + fs.StringVar(&cfg.FilePath, "file", "stores.dump", "dump file path and name") + + return cfg +} diff --git a/tools/stores-dump/main.go b/tools/pd-dev/stores-dump/main.go similarity index 57% rename from tools/stores-dump/main.go rename to tools/pd-dev/stores-dump/main.go index b3d9b00be9f..e313028d5ca 100644 --- a/tools/stores-dump/main.go +++ b/tools/pd-dev/stores-dump/main.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package storesdump import ( "bufio" @@ -27,18 +27,11 @@ import ( "github.com/pingcap/errors" "github.com/pingcap/kvproto/pkg/metapb" - "github.com/tikv/pd/pkg/utils/etcdutil" + "github.com/pingcap/log" + "github.com/tikv/pd/tools/pd-dev/util" "go.etcd.io/etcd/clientv3" "go.etcd.io/etcd/pkg/transport" -) - -var ( - clusterID = flag.Uint64("cluster-id", 0, "please make cluster ID match with TiKV") - endpoints = flag.String("endpoints", "http://127.0.0.1:2379", "endpoints urls") - filePath = flag.String("file", "stores.dump", "dump file path and name") - caPath = flag.String("cacert", "", "path of file that contains list of trusted SSL CAs") - certPath = flag.String("cert", "", "path of file that contains X509 certificate in PEM format") - keyPath = flag.String("key", "", "path of file that contains X509 key in PEM format") + "go.uber.org/zap" ) const ( @@ -59,11 +52,19 @@ func checkErr(err error) { } } -func main() { - flag.Parse() +func Run() { + cfg := newConfig() + err := cfg.Parse(os.Args[1:]) + switch errors.Cause(err) { + case nil: + case flag.ErrHelp: + os.Exit(0) + default: + log.Fatal("parse cmd flags error", zap.Error(err)) + } - rootPath = path.Join(pdRootPath, strconv.FormatUint(*clusterID, 10)) - f, err := os.Create(*filePath) + rootPath = path.Join(pdRootPath, strconv.FormatUint(cfg.ClusterID, 10)) + f, err := os.Create(cfg.FilePath) checkErr(err) defer func() { if err := f.Close(); err != nil { @@ -71,12 +72,12 @@ func main() { } }() - urls := strings.Split(*endpoints, ",") + urls := strings.Split(cfg.PDAddrs, ",") tlsInfo := transport.TLSInfo{ - CertFile: *certPath, - KeyFile: *keyPath, - TrustedCAFile: *caPath, + CertFile: cfg.CertPath, + KeyFile: cfg.KeyPath, + TrustedCAFile: cfg.CaPath, } tlsConfig, err := tlsInfo.ClientConfig() checkErr(err) @@ -100,7 +101,7 @@ func loadStores(client *clientv3.Client, f *os.File) error { defer w.Flush() for { key := path.Join(clusterPath, "s", fmt.Sprintf("%020d", nextID)) - _, res, err := loadRange(client, key, endKey, minKVRangeLimit) + _, res, err := util.LoadRange(client, rootPath, key, endKey, minKVRangeLimit) if err != nil { return err } @@ -118,22 +119,3 @@ func loadStores(client *clientv3.Client, f *os.File) error { } } } - -func loadRange(client *clientv3.Client, key, endKey string, limit int) ([]string, []string, error) { - key = path.Join(rootPath, key) - endKey = path.Join(rootPath, endKey) - - withRange := clientv3.WithRange(endKey) - withLimit := clientv3.WithLimit(int64(limit)) - resp, err := etcdutil.EtcdKVGet(client, key, withRange, withLimit) - if err != nil { - return nil, nil, err - } - keys := make([]string, 0, len(resp.Kvs)) - values := make([]string, 0, len(resp.Kvs)) - for _, item := range resp.Kvs { - keys = append(keys, strings.TrimPrefix(strings.TrimPrefix(string(item.Key), rootPath), "/")) - values = append(values, string(item.Value)) - } - return keys, values, nil -} diff --git a/tools/pd-api-bench/config/config.go b/tools/pd-dev/util/general_config.go similarity index 51% rename from tools/pd-api-bench/config/config.go rename to tools/pd-dev/util/general_config.go index d1048c0da72..f268afc8e28 100644 --- a/tools/pd-api-bench/config/config.go +++ b/tools/pd-dev/util/general_config.go @@ -12,51 +12,48 @@ // See the License for the specific language governing permissions and // limitations under the License. -package config +package util import ( + "crypto/tls" + "os" + + "github.com/pingcap/errors" "github.com/pingcap/log" - "github.com/pkg/errors" flag "github.com/spf13/pflag" + "github.com/tikv/pd/client/tlsutil" "github.com/tikv/pd/pkg/utils/configutil" - "github.com/tikv/pd/tools/pd-api-bench/cases" + "github.com/tikv/pd/pkg/utils/logutil" "go.uber.org/zap" ) -// Config is the heartbeat-bench configuration. -type Config struct { - flagSet *flag.FlagSet - configFile string - PDAddr string `toml:"pd" json:"pd"` +var defaultLogFormat = "text" + +type GeneralConfig struct { + FlagSet *flag.FlagSet + ConfigFile string + PDAddrs string `toml:"pd-endpoints" json:"pd"` StatusAddr string `toml:"status" json:"status"` Log log.Config `toml:"log" json:"log"` Logger *zap.Logger LogProps *log.ZapProperties - Client int64 `toml:"client" json:"client"` - // tls CaPath string `toml:"ca-path" json:"ca-path"` CertPath string `toml:"cert-path" json:"cert-path"` KeyPath string `toml:"key-path" json:"key-path"` - - // only for init - HTTP map[string]cases.Config `toml:"http" json:"http"` - GRPC map[string]cases.Config `toml:"grpc" json:"grpc"` - ETCD map[string]cases.Config `toml:"etcd" json:"etcd"` } -// NewConfig return a set of settings. -func NewConfig(flagSet *flag.FlagSet) *Config { - cfg := &Config{} - cfg.flagSet = flagSet - fs := cfg.flagSet - fs.StringVar(&cfg.configFile, "config", "", "config file") - fs.StringVar(&cfg.PDAddr, "pd", "http://127.0.0.1:2379", "pd address") +// NewGeneralConfig return a set of settings. +func NewGeneralConfig(flagSet *flag.FlagSet) *GeneralConfig { + cfg := &GeneralConfig{} + cfg.FlagSet = flagSet + fs := cfg.FlagSet + fs.StringVar(&cfg.ConfigFile, "config", "", "config file") + fs.StringVar(&cfg.PDAddrs, "pd-endpoints", "http://127.0.0.1:2379", "pd address") fs.StringVar(&cfg.Log.File.Filename, "log-file", "", "log file path") fs.StringVar(&cfg.StatusAddr, "status", "127.0.0.1:10081", "status address") - fs.Int64Var(&cfg.Client, "client", 1, "client number") fs.StringVar(&cfg.CaPath, "cacert", "", "path of file that contains list of trusted SSL CAs") fs.StringVar(&cfg.CertPath, "cert", "", "path of file that contains X509 certificate in PEM format") fs.StringVar(&cfg.KeyPath, "key", "", "path of file that contains X509 key in PEM format") @@ -64,60 +61,74 @@ func NewConfig(flagSet *flag.FlagSet) *Config { } // Parse parses flag definitions from the argument list. -func (c *Config) Parse(arguments []string) error { +func (c *GeneralConfig) Parse(arguments []string) error { // Parse first to get config file. - err := c.flagSet.Parse(arguments) + err := c.FlagSet.Parse(arguments) if err != nil { return errors.WithStack(err) } // Load config file if specified. - if c.configFile != "" { - _, err = configutil.ConfigFromFile(c, c.configFile) + if c.ConfigFile != "" { + _, err = configutil.ConfigFromFile(c, c.ConfigFile) if err != nil { return err } } - c.Adjust() // Parse again to replace with command line options. - err = c.flagSet.Parse(arguments) + err = c.FlagSet.Parse(arguments) if err != nil { return errors.WithStack(err) } - if len(c.flagSet.Args()) != 0 { - return errors.Errorf("'%s' is an invalid flag", c.flagSet.Arg(0)) + if len(c.FlagSet.Args()) != 0 { + return errors.Errorf("'%s' is an invalid flag", c.FlagSet.Arg(0)) } + c.Adjust() return nil } -// InitCoordinator set case config from config itself. -func (c *Config) InitCoordinator(co *cases.Coordinator) { - for name, cfg := range c.HTTP { - err := co.SetHTTPCase(name, &cfg) - if err != nil { - log.Error("create HTTP case failed", zap.Error(err)) - } - } - for name, cfg := range c.GRPC { - err := co.SetGRPCCase(name, &cfg) - if err != nil { - log.Error("create gRPC case failed", zap.Error(err)) - } +// Adjust is used to adjust configurations +func (c *GeneralConfig) Adjust() { + if len(c.Log.Format) == 0 { + c.Log.Format = defaultLogFormat } - for name, cfg := range c.ETCD { - err := co.SetETCDCase(name, &cfg) - if err != nil { - log.Error("create etcd case failed", zap.Error(err)) - } + + err := logutil.SetupLogger(c.Log, &c.Logger, &c.LogProps) + if err == nil { + log.ReplaceGlobals(c.Logger, c.LogProps) + } else { + log.Fatal("initialize logger error", zap.Error(err)) } } -// Adjust is used to adjust configurations -func (c *Config) Adjust() { - if len(c.Log.Format) == 0 { - c.Log.Format = "text" +func LoadTLSConfig(cfg *GeneralConfig) *tls.Config { + if len(cfg.CaPath) == 0 { + return nil + } + caData, err := os.ReadFile(cfg.CaPath) + if err != nil { + log.Error("fail to read ca file", zap.Error(err)) + } + certData, err := os.ReadFile(cfg.CertPath) + if err != nil { + log.Error("fail to read cert file", zap.Error(err)) + } + keyData, err := os.ReadFile(cfg.KeyPath) + if err != nil { + log.Error("fail to read key file", zap.Error(err)) } + + tlsConf, err := tlsutil.TLSConfig{ + SSLCABytes: caData, + SSLCertBytes: certData, + SSLKEYBytes: keyData, + }.ToTLSConfig() + if err != nil { + log.Fatal("failed to load tlc config", zap.Error(err)) + } + + return tlsConf } diff --git a/tools/pd-dev/util/util.go b/tools/pd-dev/util/util.go new file mode 100644 index 00000000000..5f8aec7b818 --- /dev/null +++ b/tools/pd-dev/util/util.go @@ -0,0 +1,43 @@ +// Copyright 2024 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 util + +import ( + "path" + "strings" + + "github.com/tikv/pd/pkg/utils/etcdutil" + "go.etcd.io/etcd/clientv3" +) + +func LoadRange(client *clientv3.Client, rootPath string, + key, endKey string, limit int) ([]string, []string, error) { + key = path.Join(rootPath, key) + endKey = path.Join(rootPath, endKey) + + withRange := clientv3.WithRange(endKey) + withLimit := clientv3.WithLimit(int64(limit)) + resp, err := etcdutil.EtcdKVGet(client, key, withRange, withLimit) + if err != nil { + return nil, nil, err + } + keys := make([]string, 0, len(resp.Kvs)) + values := make([]string, 0, len(resp.Kvs)) + for _, item := range resp.Kvs { + keys = append(keys, strings.TrimPrefix(strings.TrimPrefix(string(item.Key), rootPath), "/")) + values = append(values, string(item.Value)) + } + return keys, values, nil +} diff --git a/tools/pd-heartbeat-bench/main.go b/tools/pd-heartbeat-bench/main.go deleted file mode 100644 index ec5e2506e6b..00000000000 --- a/tools/pd-heartbeat-bench/main.go +++ /dev/null @@ -1,722 +0,0 @@ -// Copyright 2019 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 main - -import ( - "context" - "crypto/tls" - "fmt" - "io" - "math/rand" - "net/http" - "os" - "os/signal" - "sync" - "sync/atomic" - "syscall" - "time" - - "github.com/docker/go-units" - "github.com/gin-contrib/cors" - "github.com/gin-contrib/gzip" - "github.com/gin-contrib/pprof" - "github.com/gin-gonic/gin" - "github.com/pingcap/errors" - "github.com/pingcap/kvproto/pkg/metapb" - "github.com/pingcap/kvproto/pkg/pdpb" - "github.com/pingcap/log" - "github.com/spf13/pflag" - "github.com/tikv/pd/client/grpcutil" - pdHttp "github.com/tikv/pd/client/http" - "github.com/tikv/pd/client/tlsutil" - "github.com/tikv/pd/pkg/codec" - "github.com/tikv/pd/pkg/mcs/utils" - "github.com/tikv/pd/pkg/statistics" - "github.com/tikv/pd/pkg/utils/logutil" - "github.com/tikv/pd/tools/pd-heartbeat-bench/config" - "go.etcd.io/etcd/pkg/report" - "go.uber.org/zap" -) - -const ( - bytesUnit = 128 - keysUint = 8 - queryUnit = 8 - hotByteUnit = 16 * units.KiB - hotKeysUint = 256 - hotQueryUnit = 256 - regionReportInterval = 60 // 60s - storeReportInterval = 10 // 10s - capacity = 4 * units.TiB -) - -var ( - clusterID uint64 - maxVersion uint64 = 1 -) - -func newClient(ctx context.Context, cfg *config.Config) (pdpb.PDClient, error) { - tlsConfig, err := cfg.Security.ToTLSConfig() - if err != nil { - return nil, err - } - cc, err := grpcutil.GetClientConn(ctx, cfg.PDAddr, tlsConfig) - if err != nil { - return nil, err - } - return pdpb.NewPDClient(cc), nil -} - -func initClusterID(ctx context.Context, cli pdpb.PDClient) { - ticker := time.NewTicker(time.Second) - defer ticker.Stop() - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - cctx, cancel := context.WithCancel(ctx) - res, err := cli.GetMembers(cctx, &pdpb.GetMembersRequest{}) - cancel() - if err != nil { - continue - } - if res.GetHeader().GetError() != nil { - continue - } - clusterID = res.GetHeader().GetClusterId() - log.Info("init cluster ID successfully", zap.Uint64("cluster-id", clusterID)) - return - } - } -} - -func header() *pdpb.RequestHeader { - return &pdpb.RequestHeader{ - ClusterId: clusterID, - } -} - -func bootstrap(ctx context.Context, cli pdpb.PDClient) { - cctx, cancel := context.WithCancel(ctx) - isBootstrapped, err := cli.IsBootstrapped(cctx, &pdpb.IsBootstrappedRequest{Header: header()}) - cancel() - if err != nil { - log.Fatal("check if cluster has already bootstrapped failed", zap.Error(err)) - } - if isBootstrapped.GetBootstrapped() { - log.Info("already bootstrapped") - return - } - - store := &metapb.Store{ - Id: 1, - Address: fmt.Sprintf("localhost:%d", 2), - Version: "6.4.0-alpha", - } - region := &metapb.Region{ - Id: 1, - Peers: []*metapb.Peer{{StoreId: 1, Id: 1}}, - RegionEpoch: &metapb.RegionEpoch{ConfVer: 1, Version: 1}, - } - req := &pdpb.BootstrapRequest{ - Header: header(), - Store: store, - Region: region, - } - cctx, cancel = context.WithCancel(ctx) - resp, err := cli.Bootstrap(cctx, req) - cancel() - if err != nil { - log.Fatal("failed to bootstrap the cluster", zap.Error(err)) - } - if resp.GetHeader().GetError() != nil { - log.Fatal("failed to bootstrap the cluster", zap.String("err", resp.GetHeader().GetError().String())) - } - log.Info("bootstrapped") -} - -func putStores(ctx context.Context, cfg *config.Config, cli pdpb.PDClient, stores *Stores) { - for i := uint64(1); i <= uint64(cfg.StoreCount); i++ { - store := &metapb.Store{ - Id: i, - Address: fmt.Sprintf("localhost:%d", i), - Version: "6.4.0-alpha", - } - cctx, cancel := context.WithCancel(ctx) - resp, err := cli.PutStore(cctx, &pdpb.PutStoreRequest{Header: header(), Store: store}) - cancel() - if err != nil { - log.Fatal("failed to put store", zap.Uint64("store-id", i), zap.Error(err)) - } - if resp.GetHeader().GetError() != nil { - log.Fatal("failed to put store", zap.Uint64("store-id", i), zap.String("err", resp.GetHeader().GetError().String())) - } - go func(ctx context.Context, storeID uint64) { - var heartbeatTicker = time.NewTicker(10 * time.Second) - defer heartbeatTicker.Stop() - for { - select { - case <-heartbeatTicker.C: - stores.heartbeat(ctx, cli, storeID) - case <-ctx.Done(): - return - } - } - }(ctx, i) - } -} - -// Regions simulates all regions to heartbeat. -type Regions struct { - regions []*pdpb.RegionHeartbeatRequest - awakenRegions atomic.Value - - updateRound int - - updateLeader []int - updateEpoch []int - updateSpace []int - updateFlow []int -} - -func (rs *Regions) init(cfg *config.Config) { - rs.regions = make([]*pdpb.RegionHeartbeatRequest, 0, cfg.RegionCount) - rs.updateRound = 0 - - // Generate regions - id := uint64(1) - now := uint64(time.Now().Unix()) - - for i := 0; i < cfg.RegionCount; i++ { - region := &pdpb.RegionHeartbeatRequest{ - Header: header(), - Region: &metapb.Region{ - Id: id, - StartKey: codec.GenerateTableKey(int64(i)), - EndKey: codec.GenerateTableKey(int64(i + 1)), - RegionEpoch: &metapb.RegionEpoch{ConfVer: 2, Version: maxVersion}, - }, - ApproximateSize: bytesUnit, - Interval: &pdpb.TimeInterval{ - StartTimestamp: now, - EndTimestamp: now + regionReportInterval, - }, - QueryStats: &pdpb.QueryStats{}, - ApproximateKeys: keysUint, - Term: 1, - } - id += 1 - if i == 0 { - region.Region.StartKey = []byte("") - } - if i == cfg.RegionCount-1 { - region.Region.EndKey = []byte("") - } - - peers := make([]*metapb.Peer, 0, cfg.Replica) - for j := 0; j < cfg.Replica; j++ { - peers = append(peers, &metapb.Peer{Id: id, StoreId: uint64((i+j)%cfg.StoreCount + 1)}) - id += 1 - } - - region.Region.Peers = peers - region.Leader = peers[0] - rs.regions = append(rs.regions, region) - } -} - -func (rs *Regions) update(cfg *config.Config, options *config.Options) { - rs.updateRound += 1 - - // Generate sample index - indexes := make([]int, cfg.RegionCount) - for i := range indexes { - indexes[i] = i - } - reportRegions := pick(indexes, cfg.RegionCount, options.GetReportRatio()) - - reportCount := len(reportRegions) - rs.updateFlow = pick(reportRegions, reportCount, options.GetFlowUpdateRatio()) - rs.updateLeader = randomPick(reportRegions, reportCount, options.GetLeaderUpdateRatio()) - rs.updateEpoch = randomPick(reportRegions, reportCount, options.GetEpochUpdateRatio()) - rs.updateSpace = randomPick(reportRegions, reportCount, options.GetSpaceUpdateRatio()) - var ( - updatedStatisticsMap = make(map[int]*pdpb.RegionHeartbeatRequest) - awakenRegions []*pdpb.RegionHeartbeatRequest - ) - - // update leader - for _, i := range rs.updateLeader { - region := rs.regions[i] - region.Leader = region.Region.Peers[rs.updateRound%cfg.Replica] - } - // update epoch - for _, i := range rs.updateEpoch { - region := rs.regions[i] - region.Region.RegionEpoch.Version += 1 - if region.Region.RegionEpoch.Version > maxVersion { - maxVersion = region.Region.RegionEpoch.Version - } - } - // update space - for _, i := range rs.updateSpace { - region := rs.regions[i] - region.ApproximateSize = uint64(bytesUnit * rand.Float64()) - region.ApproximateKeys = uint64(keysUint * rand.Float64()) - } - // update flow - for _, i := range rs.updateFlow { - region := rs.regions[i] - if region.Leader.StoreId <= uint64(options.GetHotStoreCount()) { - region.BytesWritten = uint64(hotByteUnit * (1 + rand.Float64()) * 60) - region.BytesRead = uint64(hotByteUnit * (1 + rand.Float64()) * 10) - region.KeysWritten = uint64(hotKeysUint * (1 + rand.Float64()) * 60) - region.KeysRead = uint64(hotKeysUint * (1 + rand.Float64()) * 10) - region.QueryStats = &pdpb.QueryStats{ - Get: uint64(hotQueryUnit * (1 + rand.Float64()) * 10), - Put: uint64(hotQueryUnit * (1 + rand.Float64()) * 60), - } - } else { - region.BytesWritten = uint64(bytesUnit * rand.Float64()) - region.BytesRead = uint64(bytesUnit * rand.Float64()) - region.KeysWritten = uint64(keysUint * rand.Float64()) - region.KeysRead = uint64(keysUint * rand.Float64()) - region.QueryStats = &pdpb.QueryStats{ - Get: uint64(queryUnit * rand.Float64()), - Put: uint64(queryUnit * rand.Float64()), - } - } - updatedStatisticsMap[i] = region - } - // update interval - for _, region := range rs.regions { - region.Interval.StartTimestamp = region.Interval.EndTimestamp - region.Interval.EndTimestamp = region.Interval.StartTimestamp + regionReportInterval - } - for _, i := range reportRegions { - region := rs.regions[i] - // reset the statistics of the region which is not updated - if _, exist := updatedStatisticsMap[i]; !exist { - region.BytesWritten = 0 - region.BytesRead = 0 - region.KeysWritten = 0 - region.KeysRead = 0 - region.QueryStats = &pdpb.QueryStats{} - } - awakenRegions = append(awakenRegions, region) - } - - rs.awakenRegions.Store(awakenRegions) -} - -func createHeartbeatStream(ctx context.Context, cfg *config.Config) (pdpb.PDClient, pdpb.PD_RegionHeartbeatClient) { - cli, err := newClient(ctx, cfg) - if err != nil { - log.Fatal("create client error", zap.Error(err)) - } - stream, err := cli.RegionHeartbeat(ctx) - if err != nil { - log.Fatal("create stream error", zap.Error(err)) - } - - go func() { - // do nothing - for { - stream.Recv() - } - }() - return cli, stream -} - -func (rs *Regions) handleRegionHeartbeat(wg *sync.WaitGroup, stream pdpb.PD_RegionHeartbeatClient, storeID uint64, rep report.Report) { - defer wg.Done() - var regions, toUpdate []*pdpb.RegionHeartbeatRequest - updatedRegions := rs.awakenRegions.Load() - if updatedRegions == nil { - toUpdate = rs.regions - } else { - toUpdate = updatedRegions.([]*pdpb.RegionHeartbeatRequest) - } - for _, region := range toUpdate { - if region.Leader.StoreId != storeID { - continue - } - regions = append(regions, region) - } - - start := time.Now() - var err error - for _, region := range regions { - err = stream.Send(region) - rep.Results() <- report.Result{Start: start, End: time.Now(), Err: err} - if err == io.EOF { - log.Error("receive eof error", zap.Uint64("store-id", storeID), zap.Error(err)) - err := stream.CloseSend() - if err != nil { - log.Error("fail to close stream", zap.Uint64("store-id", storeID), zap.Error(err)) - } - return - } - if err != nil { - log.Error("send result error", zap.Uint64("store-id", storeID), zap.Error(err)) - return - } - } - log.Info("store finish one round region heartbeat", zap.Uint64("store-id", storeID), zap.Duration("cost-time", time.Since(start)), zap.Int("reported-region-count", len(regions))) -} - -// Stores contains store stats with lock. -type Stores struct { - stat []atomic.Value -} - -func newStores(storeCount int) *Stores { - return &Stores{ - stat: make([]atomic.Value, storeCount+1), - } -} - -func (s *Stores) heartbeat(ctx context.Context, cli pdpb.PDClient, storeID uint64) { - cctx, cancel := context.WithCancel(ctx) - defer cancel() - cli.StoreHeartbeat(cctx, &pdpb.StoreHeartbeatRequest{Header: header(), Stats: s.stat[storeID].Load().(*pdpb.StoreStats)}) -} - -func (s *Stores) update(rs *Regions) { - stats := make([]*pdpb.StoreStats, len(s.stat)) - now := uint64(time.Now().Unix()) - for i := range stats { - stats[i] = &pdpb.StoreStats{ - StoreId: uint64(i), - Capacity: capacity, - Available: capacity, - QueryStats: &pdpb.QueryStats{}, - PeerStats: make([]*pdpb.PeerStat, 0), - Interval: &pdpb.TimeInterval{ - StartTimestamp: now - storeReportInterval, - EndTimestamp: now, - }, - } - } - var toUpdate []*pdpb.RegionHeartbeatRequest - updatedRegions := rs.awakenRegions.Load() - if updatedRegions == nil { - toUpdate = rs.regions - } else { - toUpdate = updatedRegions.([]*pdpb.RegionHeartbeatRequest) - } - for _, region := range toUpdate { - for _, peer := range region.Region.Peers { - store := stats[peer.StoreId] - store.UsedSize += region.ApproximateSize - store.Available -= region.ApproximateSize - store.RegionCount += 1 - } - store := stats[region.Leader.StoreId] - if region.BytesWritten != 0 { - store.BytesWritten += region.BytesWritten - store.BytesRead += region.BytesRead - store.KeysWritten += region.KeysWritten - store.KeysRead += region.KeysRead - store.QueryStats.Get += region.QueryStats.Get - store.QueryStats.Put += region.QueryStats.Put - store.PeerStats = append(store.PeerStats, &pdpb.PeerStat{ - RegionId: region.Region.Id, - ReadKeys: region.KeysRead, - ReadBytes: region.BytesRead, - WrittenKeys: region.KeysWritten, - WrittenBytes: region.BytesWritten, - QueryStats: region.QueryStats, - }) - } - } - for i := range stats { - s.stat[i].Store(stats[i]) - } -} - -func randomPick(slice []int, total int, ratio float64) []int { - rand.Shuffle(total, func(i, j int) { - slice[i], slice[j] = slice[j], slice[i] - }) - return append(slice[:0:0], slice[0:int(float64(total)*ratio)]...) -} - -func pick(slice []int, total int, ratio float64) []int { - return append(slice[:0:0], slice[0:int(float64(total)*ratio)]...) -} - -func main() { - rand.New(rand.NewSource(0)) // Ensure consistent behavior multiple times - statistics.Denoising = false - cfg := config.NewConfig() - err := cfg.Parse(os.Args[1:]) - defer logutil.LogPanic() - - switch errors.Cause(err) { - case nil: - case pflag.ErrHelp: - exit(0) - default: - log.Fatal("parse cmd flags error", zap.Error(err)) - } - - // New zap logger - err = logutil.SetupLogger(cfg.Log, &cfg.Logger, &cfg.LogProps) - if err == nil { - log.ReplaceGlobals(cfg.Logger, cfg.LogProps) - } else { - log.Fatal("initialize logger error", zap.Error(err)) - } - - maxVersion = cfg.InitEpochVer - options := config.NewOptions(cfg) - // let PD have enough time to start - time.Sleep(5 * time.Second) - ctx, cancel := context.WithCancel(context.Background()) - sc := make(chan os.Signal, 1) - signal.Notify(sc, - syscall.SIGHUP, - syscall.SIGINT, - syscall.SIGTERM, - syscall.SIGQUIT) - - var sig os.Signal - go func() { - sig = <-sc - cancel() - }() - cli, err := newClient(ctx, cfg) - if err != nil { - log.Fatal("create client error", zap.Error(err)) - } - - initClusterID(ctx, cli) - go runHTTPServer(cfg, options) - regions := new(Regions) - regions.init(cfg) - log.Info("finish init regions") - stores := newStores(cfg.StoreCount) - stores.update(regions) - bootstrap(ctx, cli) - putStores(ctx, cfg, cli, stores) - log.Info("finish put stores") - clis := make(map[uint64]pdpb.PDClient, cfg.StoreCount) - httpCli := pdHttp.NewClient("tools-heartbeat-bench", []string{cfg.PDAddr}, pdHttp.WithTLSConfig(loadTLSConfig(cfg))) - go deleteOperators(ctx, httpCli) - streams := make(map[uint64]pdpb.PD_RegionHeartbeatClient, cfg.StoreCount) - for i := 1; i <= cfg.StoreCount; i++ { - clis[uint64(i)], streams[uint64(i)] = createHeartbeatStream(ctx, cfg) - } - header := &pdpb.RequestHeader{ - ClusterId: clusterID, - } - var heartbeatTicker = time.NewTicker(regionReportInterval * time.Second) - defer heartbeatTicker.Stop() - var resolvedTSTicker = time.NewTicker(time.Second) - defer resolvedTSTicker.Stop() - for { - select { - case <-heartbeatTicker.C: - if cfg.Round != 0 && regions.updateRound > cfg.Round { - exit(0) - } - rep := newReport(cfg) - r := rep.Stats() - - startTime := time.Now() - wg := &sync.WaitGroup{} - for i := 1; i <= cfg.StoreCount; i++ { - id := uint64(i) - wg.Add(1) - go regions.handleRegionHeartbeat(wg, streams[id], id, rep) - } - wg.Wait() - - since := time.Since(startTime).Seconds() - close(rep.Results()) - regions.result(cfg.RegionCount, since) - stats := <-r - log.Info("region heartbeat stats", zap.String("total", fmt.Sprintf("%.4fs", stats.Total.Seconds())), - zap.String("slowest", fmt.Sprintf("%.4fs", stats.Slowest)), - zap.String("fastest", fmt.Sprintf("%.4fs", stats.Fastest)), - zap.String("average", fmt.Sprintf("%.4fs", stats.Average)), - zap.String("stddev", fmt.Sprintf("%.4fs", stats.Stddev)), - zap.String("rps", fmt.Sprintf("%.4f", stats.RPS)), - zap.Uint64("max-epoch-version", maxVersion), - ) - log.Info("store heartbeat stats", zap.String("max", fmt.Sprintf("%.4fs", since))) - regions.update(cfg, options) - go stores.update(regions) // update stores in background, unusually region heartbeat is slower than store update. - case <-resolvedTSTicker.C: - wg := &sync.WaitGroup{} - for i := 1; i <= cfg.StoreCount; i++ { - id := uint64(i) - wg.Add(1) - go func(wg *sync.WaitGroup, id uint64) { - defer wg.Done() - cli := clis[id] - _, err := cli.ReportMinResolvedTS(ctx, &pdpb.ReportMinResolvedTsRequest{ - Header: header, - StoreId: id, - MinResolvedTs: uint64(time.Now().Unix()), - }) - if err != nil { - log.Error("send resolved TS error", zap.Uint64("store-id", id), zap.Error(err)) - return - } - }(wg, id) - } - wg.Wait() - case <-ctx.Done(): - log.Info("got signal to exit") - switch sig { - case syscall.SIGTERM: - exit(0) - default: - exit(1) - } - } - } -} - -func exit(code int) { - os.Exit(code) -} - -func deleteOperators(ctx context.Context, httpCli pdHttp.Client) { - ticker := time.NewTicker(30 * time.Second) - defer ticker.Stop() - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - err := httpCli.DeleteOperators(ctx) - if err != nil { - log.Error("fail to delete operators", zap.Error(err)) - } - } - } -} - -func newReport(cfg *config.Config) report.Report { - p := "%4.4f" - if cfg.Sample { - return report.NewReportSample(p) - } - return report.NewReport(p) -} - -func (rs *Regions) result(regionCount int, sec float64) { - if rs.updateRound == 0 { - // There was no difference in the first round - return - } - - updated := make(map[int]struct{}) - for _, i := range rs.updateLeader { - updated[i] = struct{}{} - } - for _, i := range rs.updateEpoch { - updated[i] = struct{}{} - } - for _, i := range rs.updateSpace { - updated[i] = struct{}{} - } - for _, i := range rs.updateFlow { - updated[i] = struct{}{} - } - inactiveCount := regionCount - len(updated) - - log.Info("update speed of each category", zap.String("rps", fmt.Sprintf("%.4f", float64(regionCount)/sec)), - zap.String("save-tree", fmt.Sprintf("%.4f", float64(len(rs.updateLeader))/sec)), - zap.String("save-kv", fmt.Sprintf("%.4f", float64(len(rs.updateEpoch))/sec)), - zap.String("save-space", fmt.Sprintf("%.4f", float64(len(rs.updateSpace))/sec)), - zap.String("save-flow", fmt.Sprintf("%.4f", float64(len(rs.updateFlow))/sec)), - zap.String("skip", fmt.Sprintf("%.4f", float64(inactiveCount)/sec))) -} - -func runHTTPServer(cfg *config.Config, options *config.Options) { - gin.SetMode(gin.ReleaseMode) - engine := gin.New() - engine.Use(gin.Recovery()) - engine.Use(cors.Default()) - engine.Use(gzip.Gzip(gzip.DefaultCompression)) - engine.GET("metrics", utils.PromHandler()) - // profile API - pprof.Register(engine) - engine.PUT("config", func(c *gin.Context) { - newCfg := cfg.Clone() - newCfg.HotStoreCount = options.GetHotStoreCount() - newCfg.FlowUpdateRatio = options.GetFlowUpdateRatio() - newCfg.LeaderUpdateRatio = options.GetLeaderUpdateRatio() - newCfg.EpochUpdateRatio = options.GetEpochUpdateRatio() - newCfg.SpaceUpdateRatio = options.GetSpaceUpdateRatio() - newCfg.ReportRatio = options.GetReportRatio() - if err := c.BindJSON(&newCfg); err != nil { - c.String(http.StatusBadRequest, err.Error()) - return - } - if err := newCfg.Validate(); err != nil { - c.String(http.StatusBadRequest, err.Error()) - return - } - options.SetOptions(newCfg) - c.String(http.StatusOK, "Successfully updated the configuration") - }) - engine.GET("config", func(c *gin.Context) { - output := cfg.Clone() - output.HotStoreCount = options.GetHotStoreCount() - output.FlowUpdateRatio = options.GetFlowUpdateRatio() - output.LeaderUpdateRatio = options.GetLeaderUpdateRatio() - output.EpochUpdateRatio = options.GetEpochUpdateRatio() - output.SpaceUpdateRatio = options.GetSpaceUpdateRatio() - output.ReportRatio = options.GetReportRatio() - - c.IndentedJSON(http.StatusOK, output) - }) - engine.Run(cfg.StatusAddr) -} - -func loadTLSConfig(cfg *config.Config) *tls.Config { - if len(cfg.Security.CAPath) == 0 { - return nil - } - caData, err := os.ReadFile(cfg.Security.CAPath) - if err != nil { - log.Error("fail to read ca file", zap.Error(err)) - } - certData, err := os.ReadFile(cfg.Security.CertPath) - if err != nil { - log.Error("fail to read cert file", zap.Error(err)) - } - keyData, err := os.ReadFile(cfg.Security.KeyPath) - if err != nil { - log.Error("fail to read key file", zap.Error(err)) - } - - tlsConf, err := tlsutil.TLSConfig{ - SSLCABytes: caData, - SSLCertBytes: certData, - SSLKEYBytes: keyData, - }.ToTLSConfig() - if err != nil { - log.Fatal("failed to load tlc config", zap.Error(err)) - } - - return tlsConf -} diff --git a/tools/pd-tso-bench/main.go b/tools/pd-tso-bench/main.go deleted file mode 100644 index b4101bda270..00000000000 --- a/tools/pd-tso-bench/main.go +++ /dev/null @@ -1,459 +0,0 @@ -// Copyright 2017 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 main - -import ( - "context" - "flag" - "fmt" - "io" - "math/rand" - "net/http" - "net/http/httptest" - "os" - "os/signal" - "sync" - "syscall" - "time" - - "github.com/influxdata/tdigest" - "github.com/pingcap/errors" - "github.com/pingcap/log" - "github.com/prometheus/client_golang/prometheus/promhttp" - pd "github.com/tikv/pd/client" - "go.uber.org/zap" - "google.golang.org/grpc" - "google.golang.org/grpc/keepalive" -) - -const ( - keepaliveTime = 10 * time.Second - keepaliveTimeout = 3 * time.Second -) - -var ( - pdAddrs = flag.String("pd", "127.0.0.1:2379", "pd address") - clientNumber = flag.Int("client", 1, "the number of pd clients involved in each benchmark") - concurrency = flag.Int("c", 1000, "concurrency") - count = flag.Int("count", 1, "the count number that the test will run") - duration = flag.Duration("duration", 60*time.Second, "how many seconds the test will last") - dcLocation = flag.String("dc", "global", "which dc-location this bench will request") - verbose = flag.Bool("v", false, "output statistics info every interval and output metrics info at the end") - interval = flag.Duration("interval", time.Second, "interval to output the statistics") - caPath = flag.String("cacert", "", "path of file that contains list of trusted SSL CAs") - certPath = flag.String("cert", "", "path of file that contains X509 certificate in PEM format") - keyPath = flag.String("key", "", "path of file that contains X509 key in PEM format") - maxBatchWaitInterval = flag.Duration("batch-interval", 0, "the max batch wait interval") - enableTSOFollowerProxy = flag.Bool("enable-tso-follower-proxy", false, "whether enable the TSO Follower Proxy") - enableFaultInjection = flag.Bool("enable-fault-injection", false, "whether enable fault injection") - faultInjectionRate = flag.Float64("fault-injection-rate", 0.01, "the failure rate [0.0001, 1]. 0.01 means 1% failure rate") - maxTSOSendIntervalMilliseconds = flag.Int("max-send-interval-ms", 0, "max tso send interval in milliseconds, 60s by default") - keyspaceID = flag.Uint("keyspace-id", 0, "the id of the keyspace to access") - keyspaceName = flag.String("keyspace-name", "", "the name of the keyspace to access") - wg sync.WaitGroup -) - -var promServer *httptest.Server - -func collectMetrics(server *httptest.Server) string { - time.Sleep(1100 * time.Millisecond) - res, _ := http.Get(server.URL) - body, _ := io.ReadAll(res.Body) - res.Body.Close() - return string(body) -} - -func main() { - flag.Parse() - ctx, cancel := context.WithCancel(context.Background()) - - sc := make(chan os.Signal, 1) - signal.Notify(sc, - syscall.SIGHUP, - syscall.SIGINT, - syscall.SIGTERM, - syscall.SIGQUIT) - go func() { - <-sc - cancel() - }() - - for i := 0; i < *count; i++ { - fmt.Printf("\nStart benchmark #%d, duration: %+vs\n", i, duration.Seconds()) - bench(ctx) - } -} - -func bench(mainCtx context.Context) { - promServer = httptest.NewServer(promhttp.Handler()) - - // Initialize all clients - fmt.Printf("Create %d client(s) for benchmark\n", *clientNumber) - pdClients := make([]pd.Client, *clientNumber) - for idx := range pdClients { - pdCli, err := createPDClient(mainCtx) - if err != nil { - log.Fatal(fmt.Sprintf("create pd client #%d failed: %v", idx, err)) - } - pdClients[idx] = pdCli - } - - ctx, cancel := context.WithCancel(mainCtx) - // To avoid the first time high latency. - for idx, pdCli := range pdClients { - _, _, err := pdCli.GetLocalTS(ctx, *dcLocation) - if err != nil { - log.Fatal("get first time tso failed", zap.Int("client-number", idx), zap.Error(err)) - } - } - - durCh := make(chan time.Duration, 2*(*concurrency)*(*clientNumber)) - - if *enableFaultInjection { - fmt.Printf("Enable fault injection, failure rate: %f\n", *faultInjectionRate) - wg.Add(*clientNumber) - for i := 0; i < *clientNumber; i++ { - go reqWorker(ctx, pdClients, i, durCh) - } - } else { - wg.Add((*concurrency) * (*clientNumber)) - for i := 0; i < *clientNumber; i++ { - for j := 0; j < *concurrency; j++ { - go reqWorker(ctx, pdClients, i, durCh) - } - } - } - - wg.Add(1) - go showStats(ctx, durCh) - - timer := time.NewTimer(*duration) - defer timer.Stop() - - select { - case <-ctx.Done(): - case <-timer.C: - } - cancel() - - wg.Wait() - - for _, pdCli := range pdClients { - pdCli.Close() - } -} - -var latencyTDigest *tdigest.TDigest = tdigest.New() - -func showStats(ctx context.Context, durCh chan time.Duration) { - defer wg.Done() - - statCtx, cancel := context.WithCancel(ctx) - defer cancel() - - ticker := time.NewTicker(*interval) - defer ticker.Stop() - - s := newStats() - total := newStats() - - fmt.Println() - for { - select { - case <-ticker.C: - // runtime.GC() - if *verbose { - fmt.Println(s.Counter()) - } - total.merge(s) - s = newStats() - case d := <-durCh: - s.update(d) - case <-statCtx.Done(): - fmt.Println("\nTotal:") - fmt.Println(total.Counter()) - fmt.Println(total.Percentage()) - // Calculate the percentiles by using the tDigest algorithm. - fmt.Printf("P0.5: %.4fms, P0.8: %.4fms, P0.9: %.4fms, P0.99: %.4fms\n\n", latencyTDigest.Quantile(0.5), latencyTDigest.Quantile(0.8), latencyTDigest.Quantile(0.9), latencyTDigest.Quantile(0.99)) - if *verbose { - fmt.Println(collectMetrics(promServer)) - } - return - } - } -} - -const ( - twoDur = time.Millisecond * 2 - fiveDur = time.Millisecond * 5 - tenDur = time.Millisecond * 10 - thirtyDur = time.Millisecond * 30 - fiftyDur = time.Millisecond * 50 - oneHundredDur = time.Millisecond * 100 - twoHundredDur = time.Millisecond * 200 - fourHundredDur = time.Millisecond * 400 - eightHundredDur = time.Millisecond * 800 - oneThousandDur = time.Millisecond * 1000 -) - -type stats struct { - maxDur time.Duration - minDur time.Duration - totalDur time.Duration - count int - submilliCnt int - milliCnt int - twoMilliCnt int - fiveMilliCnt int - tenMSCnt int - thirtyCnt int - fiftyCnt int - oneHundredCnt int - twoHundredCnt int - fourHundredCnt int - eightHundredCnt int - oneThousandCnt int -} - -func newStats() *stats { - return &stats{ - minDur: time.Hour, - maxDur: 0, - } -} - -func (s *stats) update(dur time.Duration) { - s.count++ - s.totalDur += dur - latencyTDigest.Add(float64(dur.Nanoseconds())/1e6, 1) - - if dur > s.maxDur { - s.maxDur = dur - } - if dur < s.minDur { - s.minDur = dur - } - - if dur > oneThousandDur { - s.oneThousandCnt++ - return - } - - if dur > eightHundredDur { - s.eightHundredCnt++ - return - } - - if dur > fourHundredDur { - s.fourHundredCnt++ - return - } - - if dur > twoHundredDur { - s.twoHundredCnt++ - return - } - - if dur > oneHundredDur { - s.oneHundredCnt++ - return - } - - if dur > fiftyDur { - s.fiftyCnt++ - return - } - - if dur > thirtyDur { - s.thirtyCnt++ - return - } - - if dur > tenDur { - s.tenMSCnt++ - return - } - - if dur > fiveDur { - s.fiveMilliCnt++ - return - } - - if dur > twoDur { - s.twoMilliCnt++ - return - } - - if dur > time.Millisecond { - s.milliCnt++ - return - } - - s.submilliCnt++ -} - -func (s *stats) merge(other *stats) { - if s.maxDur < other.maxDur { - s.maxDur = other.maxDur - } - if s.minDur > other.minDur { - s.minDur = other.minDur - } - - s.count += other.count - s.totalDur += other.totalDur - s.submilliCnt += other.submilliCnt - s.milliCnt += other.milliCnt - s.twoMilliCnt += other.twoMilliCnt - s.fiveMilliCnt += other.fiveMilliCnt - s.tenMSCnt += other.tenMSCnt - s.thirtyCnt += other.thirtyCnt - s.fiftyCnt += other.fiftyCnt - s.oneHundredCnt += other.oneHundredCnt - s.twoHundredCnt += other.twoHundredCnt - s.fourHundredCnt += other.fourHundredCnt - s.eightHundredCnt += other.eightHundredCnt - s.oneThousandCnt += other.oneThousandCnt -} - -func (s *stats) Counter() string { - return fmt.Sprintf( - "count: %d, max: %.4fms, min: %.4fms, avg: %.4fms\n<1ms: %d, >1ms: %d, >2ms: %d, >5ms: %d, >10ms: %d, >30ms: %d, >50ms: %d, >100ms: %d, >200ms: %d, >400ms: %d, >800ms: %d, >1s: %d", - s.count, float64(s.maxDur.Nanoseconds())/float64(time.Millisecond), float64(s.minDur.Nanoseconds())/float64(time.Millisecond), float64(s.totalDur.Nanoseconds())/float64(s.count)/float64(time.Millisecond), - s.submilliCnt, s.milliCnt, s.twoMilliCnt, s.fiveMilliCnt, s.tenMSCnt, s.thirtyCnt, s.fiftyCnt, s.oneHundredCnt, s.twoHundredCnt, s.fourHundredCnt, - s.eightHundredCnt, s.oneThousandCnt) -} - -func (s *stats) Percentage() string { - return fmt.Sprintf( - "count: %d, <1ms: %2.2f%%, >1ms: %2.2f%%, >2ms: %2.2f%%, >5ms: %2.2f%%, >10ms: %2.2f%%, >30ms: %2.2f%%, >50ms: %2.2f%%, >100ms: %2.2f%%, >200ms: %2.2f%%, >400ms: %2.2f%%, >800ms: %2.2f%%, >1s: %2.2f%%", s.count, - s.calculate(s.submilliCnt), s.calculate(s.milliCnt), s.calculate(s.twoMilliCnt), s.calculate(s.fiveMilliCnt), s.calculate(s.tenMSCnt), s.calculate(s.thirtyCnt), s.calculate(s.fiftyCnt), - s.calculate(s.oneHundredCnt), s.calculate(s.twoHundredCnt), s.calculate(s.fourHundredCnt), s.calculate(s.eightHundredCnt), s.calculate(s.oneThousandCnt)) -} - -func (s *stats) calculate(count int) float64 { - return float64(count) * 100 / float64(s.count) -} - -func reqWorker(ctx context.Context, pdClients []pd.Client, clientIdx int, durCh chan time.Duration) { - defer wg.Done() - - reqCtx, cancel := context.WithCancel(ctx) - defer cancel() - var ( - err error - maxRetryTime int = 120 - sleepIntervalOnFailure time.Duration = 1000 * time.Millisecond - totalSleepBeforeGetTS time.Duration - ) - pdCli := pdClients[clientIdx] - - for { - if pdCli == nil || (*enableFaultInjection && shouldInjectFault()) { - if pdCli != nil { - pdCli.Close() - } - pdCli, err = createPDClient(ctx) - if err != nil { - log.Error(fmt.Sprintf("re-create pd client #%d failed: %v", clientIdx, err)) - select { - case <-reqCtx.Done(): - case <-time.After(100 * time.Millisecond): - } - continue - } - pdClients[clientIdx] = pdCli - } - - totalSleepBeforeGetTS = 0 - start := time.Now() - - i := 0 - for ; i < maxRetryTime; i++ { - var ticker *time.Ticker - if *maxTSOSendIntervalMilliseconds > 0 { - sleepBeforeGetTS := time.Duration(rand.Intn(*maxTSOSendIntervalMilliseconds)) * time.Millisecond - ticker = time.NewTicker(sleepBeforeGetTS) - select { - case <-reqCtx.Done(): - case <-ticker.C: - totalSleepBeforeGetTS += sleepBeforeGetTS - } - } - _, _, err = pdCli.GetLocalTS(reqCtx, *dcLocation) - if errors.Cause(err) == context.Canceled { - ticker.Stop() - return - } - if err == nil { - ticker.Stop() - break - } - log.Error(fmt.Sprintf("%v", err)) - time.Sleep(sleepIntervalOnFailure) - } - if err != nil { - log.Fatal(fmt.Sprintf("%v", err)) - } - dur := time.Since(start) - time.Duration(i)*sleepIntervalOnFailure - totalSleepBeforeGetTS - - select { - case <-reqCtx.Done(): - return - case durCh <- dur: - } - } -} - -func createPDClient(ctx context.Context) (pd.Client, error) { - var ( - pdCli pd.Client - err error - ) - - opts := make([]pd.ClientOption, 0) - opts = append(opts, pd.WithGRPCDialOptions( - grpc.WithKeepaliveParams(keepalive.ClientParameters{ - Time: keepaliveTime, - Timeout: keepaliveTimeout, - }), - )) - - if len(*keyspaceName) > 0 { - apiCtx := pd.NewAPIContextV2(*keyspaceName) - pdCli, err = pd.NewClientWithAPIContext(ctx, apiCtx, []string{*pdAddrs}, pd.SecurityOption{ - CAPath: *caPath, - CertPath: *certPath, - KeyPath: *keyPath, - }, opts...) - } else { - pdCli, err = pd.NewClientWithKeyspace(ctx, uint32(*keyspaceID), []string{*pdAddrs}, pd.SecurityOption{ - CAPath: *caPath, - CertPath: *certPath, - KeyPath: *keyPath, - }, opts...) - } - if err != nil { - return nil, err - } - - pdCli.UpdateOption(pd.MaxTSOBatchWaitInterval, *maxBatchWaitInterval) - pdCli.UpdateOption(pd.EnableTSOFollowerProxy, *enableTSOFollowerProxy) - return pdCli, err -} - -func shouldInjectFault() bool { - return rand.Intn(10000) < int(*faultInjectionRate*10000) -}