examples: add the Benchmark tool for TiKV

examples/benchtool/config/global.go

@@ -0,0 +1,149 @@
+// Copyright 2024 TiKV 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 (
+	"fmt"
+	"os"
+	"strconv"
+	"sync"
+	"time"
+
+	"github.com/pingcap/errors"
+	"github.com/pingcap/log"
+	"github.com/spf13/cobra"
+	"github.com/tikv/client-go/v2/config"
+	"go.uber.org/zap"
+)
+
+var initOnce = sync.Once{}
+
+const (
+	WorkloadColumnFamilyDefault = "default"
+	WorkloadColumnFamilyWrite   = "write"
+	WorkloadColumnFamilyLock    = "lock"
+)
+
+type GlobalConfig struct {
+	hosts      []string
+	port       int
+	StatusPort int
+
+	Threads        int
+	TotalTime      time.Duration
+	TotalCount     int
+	DropData       bool
+	IgnoreError    bool
+	OutputInterval time.Duration
+	Silence        bool
+	OutputStyle    string
+
+	Targets  []string
+	Security config.Security
+
+	// for log
+	LogLevel string
+	LogFile  string
+}
+
+func (c *GlobalConfig) ParsePdAddrs() error {
+	if len(c.hosts) == 0 {
+		return fmt.Errorf("PD address is empty")
+	}
+	targets := make([]string, 0, len(c.hosts))
+	for _, host := range c.hosts {
+		targets = append(targets, host+":"+strconv.Itoa(c.port))
+	}
+	c.Targets = targets
+	return nil
+}
+
+func (c *GlobalConfig) Format() string {
+	return fmt.Sprintf("Hosts: %v, Port: %d, StatusPort: %d, Threads: %d, TotalTime: %v, TotalCount: %d, DropData: %t, IgnoreError: %t, OutputInterval: %v, Silence: %t, OutputStyle: %s",
+		c.hosts, c.port, c.StatusPort, c.Threads, c.TotalTime, c.TotalCount, c.DropData, c.IgnoreError, c.OutputInterval, c.Silence, c.OutputStyle)
+}
+
+func (c *GlobalConfig) InitLogger() (err error) {
+	initOnce.Do(func() {
+		// Initialize the logger.
+		conf := &log.Config{
+			Level: c.LogLevel,
+			File: log.FileLogConfig{
+				Filename: c.LogFile,
+				MaxSize:  256,
+			},
+		}
+		lg, p, e := log.InitLogger(conf)
+		if e != nil {
+			err = e
+			return
+		}
+		log.ReplaceGlobals(lg, p)
+	})
+	return errors.Trace(err)
+}
+
+type CommandLineParser struct {
+	command *cobra.Command
+	config  *GlobalConfig
+}
+
+func NewCommandLineParser() *CommandLineParser {
+	return &CommandLineParser{}
+}
+
+func (p *CommandLineParser) Initialize() {
+	var globalCfg = &GlobalConfig{}
+	var rootCmd = &cobra.Command{
+		Use:   "bench-tool",
+		Short: "Benchmark tikv with different workloads",
+		PersistentPreRun: func(cmd *cobra.Command, args []string) {
+			if err := globalCfg.InitLogger(); err != nil {
+				log.Error("InitLogger failed", zap.Error(err))
+			}
+		},
+	}
+
+	rootCmd.PersistentFlags().StringSliceVarP(&globalCfg.hosts, "host", "H", []string{"127.0.0.1"}, "PD host")
+	rootCmd.PersistentFlags().IntVarP(&globalCfg.port, "port", "P", 4000, "PD port")
+	rootCmd.PersistentFlags().IntVarP(&globalCfg.StatusPort, "statusPort", "S", 10080, "PD status port")
+
+	rootCmd.PersistentFlags().IntVarP(&globalCfg.Threads, "threads", "T", 1, "Thread concurrency")
+	rootCmd.PersistentFlags().DurationVar(&globalCfg.TotalTime, "time", 1<<63-1, "Total execution time")
+	rootCmd.PersistentFlags().IntVar(&globalCfg.TotalCount, "count", 0, "Total execution count, 0 means infinite")
+	rootCmd.PersistentFlags().BoolVar(&globalCfg.DropData, "dropdata", false, "Cleanup data before prepare")
+	rootCmd.PersistentFlags().BoolVar(&globalCfg.IgnoreError, "ignore-error", false, "Ignore error when running workload")
+	rootCmd.PersistentFlags().BoolVar(&globalCfg.Silence, "silence", false, "Don't print error when running workload")
+	rootCmd.PersistentFlags().DurationVar(&globalCfg.OutputInterval, "interval", 10*time.Second, "Output interval time")
+	rootCmd.PersistentFlags().StringVar(&globalCfg.OutputStyle, "output", "plain", "output style, valid values can be { plain | table | json }")
+
+	rootCmd.PersistentFlags().StringVar(&globalCfg.LogFile, "log-file", "record.log", "filename of the log file")
+	rootCmd.PersistentFlags().StringVar(&globalCfg.LogLevel, "log-level", "info", "log level { debug | info | warn | error | fatal }")
+
+	rootCmd.SetOut(os.Stdout)
+
+	cobra.EnablePrefixMatching = true
+
+	p.command = rootCmd
+	p.config = globalCfg
+}
+
+func (p *CommandLineParser) GetConfig() *GlobalConfig {
+	return p.config
+}
+
+func (p *CommandLineParser) GetCommand() *cobra.Command {
+	return p.command
+}

examples/benchtool/config/pattern.go

@@ -0,0 +1,154 @@
+// Copyright 2024 TiKV 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 (
+	"fmt"
+	"os"
+
+	"gopkg.in/yaml.v3"
+)
+
+const (
+	WorkloadTypeHybrid = "hybrid"
+)
+
+// TODO: convert the txnConfig and rawkvConfig to interfaces
+type SubPatternConfig struct {
+	txnConfig   *TxnKVConfig
+	rawkvConfig *RawKVConfig
+	workloads   []string
+	name        string
+}
+
+func (s *SubPatternConfig) GetName() string {
+	return s.name
+}
+
+func (s *SubPatternConfig) GetWorkloads() []string {
+	return s.workloads
+}
+
+func (s *SubPatternConfig) GetTxnKVConfig() *TxnKVConfig {
+	return s.txnConfig
+}
+
+func (s *SubPatternConfig) GetRawKVConfig() *RawKVConfig {
+	return s.rawkvConfig
+}
+
+type SubPattern struct {
+	Name         string   `yaml:"name,omitempty"`
+	WorkloadType string   `yaml:"workload_type,omitempty"`
+	Workloads    []string `yaml:"workloads,omitempty"`
+
+	// for txnkv
+	Mode        string `yaml:"mode,omitempty"`
+	LockTimeout int    `yaml:"lock_timeout"`
+	ColumnSize  int    `yaml:"column_size"`
+	TxnSize     int    `yaml:"txn_size"`
+	// common
+	Count     int    `yaml:"count"`
+	KeyPrefix string `yaml:"key_prefix,omitempty"`
+	KeySize   int    `yaml:"key_size"`
+	ValueSize int    `yaml:"value_size"`
+	BatchSize int    `yaml:"batch_size"`
+	Threads   int    `yaml:"threads"`
+	Randomize bool   `yaml:"random"`
+}
+
+func (s *SubPattern) ConvertBasedOn(global *GlobalConfig) *SubPatternConfig {
+	// Invalid workloads
+	if s.Workloads == nil {
+		return nil
+	}
+
+	globalCfg := &GlobalConfig{}
+	if global != nil {
+		globalCfg = global
+	}
+	globalCfg.TotalCount = s.Count
+	globalCfg.Threads = s.Threads
+
+	switch s.WorkloadType {
+	case WorkloadTypeTxnKV:
+		config := &TxnKVConfig{
+			TxnMode:     s.Mode,
+			LockTimeout: s.LockTimeout,
+			// KeyPrefix:   s.key_prefix,
+			KeySize:    s.KeySize,
+			ValueSize:  s.ValueSize,
+			ColumnSize: s.ColumnSize,
+			TxnSize:    s.TxnSize,
+		}
+		config.Global = globalCfg
+		return &SubPatternConfig{
+			txnConfig: config,
+			workloads: s.Workloads,
+			name:      s.Name,
+		}
+	case WorkloadTypeRawKV:
+		config := &RawKVConfig{
+			// KeyPrefix: s.key_prefix,
+			KeySize:   s.KeySize,
+			BatchSize: s.BatchSize,
+			ValueSize: s.ValueSize,
+			Randomize: s.Randomize,
+		}
+		config.Global = globalCfg
+		return &SubPatternConfig{
+			rawkvConfig: config,
+			workloads:   s.Workloads,
+			name:        s.Name,
+		}
+	}
+	return nil
+}
+
+type PatternsConfig struct {
+	Items []*SubPattern `yaml:"patterns"`
+
+	FilePath string
+
+	Plans  []*SubPatternConfig
+	Global *GlobalConfig
+}
+
+// Parse parses the yaml file.
+func (p *PatternsConfig) Parse() error {
+	data, err := os.ReadFile(p.FilePath)
+	if err != nil {
+		return err
+	}
+	err = yaml.Unmarshal(data, p)
+	if err != nil {
+		return err
+	}
+	p.Plans = make([]*SubPatternConfig, 0, len(p.Items))
+	for _, item := range p.Items {
+		p.Plans = append(p.Plans, item.ConvertBasedOn(p.Global))
+	}
+	return nil
+}
+
+func (p *PatternsConfig) Validate() error {
+	if p.Global == nil {
+		return fmt.Errorf("global config is missing")
+	}
+	if p.Items == nil {
+		return fmt.Errorf("patterns config is missing")
+	}
+	return p.Global.ParsePdAddrs()
+}

examples/benchtool/config/rawkv.go

@@ -0,0 +1,66 @@
+// Copyright 2024 TiKV 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 (
+	"fmt"
+	"time"
+)
+
+const (
+	WorkloadTypeRawKV = "rawkv"
+)
+
+const (
+	RawKVCommandTypePut         = "put"
+	RawKVCommandTypeGet         = "get"
+	RawKVCommandTypeDel         = "del"
+	RawKVCommandTypeBatchPut    = "batch_put"
+	RawKVCommandTypeBatchGet    = "batch_get"
+	RawKVCommandTypeBatchDel    = "batch_del"
+	RawKVCommandTypeScan        = "scan"
+	RawKVCommandTypeReverseScan = "reverse_scan"
+	RawKVCommandTypeCAS         = "cas"
+
+	RawKVCommandDefaultKey    = "rawkv_key"
+	RawKVCommandDefaultEndKey = "rawkv_key`"
+	RawKVCommandDefaultValue  = "rawkv_value"
+)
+
+type RawKVConfig struct {
+	KeySize   int
+	ValueSize int
+	BatchSize int
+
+	ColumnFamily         string
+	CommandType          string
+	PrepareRetryCount    int
+	PrepareRetryInterval time.Duration
+	Randomize            bool
+
+	Global *GlobalConfig
+}
+
+func (c *RawKVConfig) Validate() error {
+	if c.KeySize <= 0 || c.ValueSize <= 0 {
+		return fmt.Errorf("key size or value size must be greater than 0")
+	}
+	if c.ColumnFamily != WorkloadColumnFamilyDefault &&
+		c.ColumnFamily != WorkloadColumnFamilyWrite &&
+		c.ColumnFamily != WorkloadColumnFamilyLock {
+		return fmt.Errorf("invalid column family: %s", c.ColumnFamily)
+	}
+	return c.Global.ParsePdAddrs()
+}