tools: rewrite benchmark tool

tools/benchmark/cmd/put.go

@@ -0,0 +1,107 @@
+// Copyright 2015 CoreOS, Inc.
+//
+// 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 cmd
+
+import (
+	"time"
+
+	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/cheggaaa/pb"
+	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/spf13/cobra"
+	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
+	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"
+	"github.com/coreos/etcd/etcdserver/etcdserverpb"
+)
+
+// putCmd represents the put command
+var putCmd = &cobra.Command{
+	Use:   "put",
+	Short: "Benchmark put",
+
+	Run: putFunc,
+}
+
+var (
+	keySize int
+	valSize int
+
+	putTotal int
+)
+
+func init() {
+	RootCmd.AddCommand(putCmd)
+	putCmd.Flags().IntVar(&keySize, "key-size", 8, "Key size of put request")
+	putCmd.Flags().IntVar(&valSize, "val-size", 8, "Value size of put request")
+	putCmd.Flags().IntVar(&putTotal, "total", 10000, "Total number of put requests")
+}
+
+func putFunc(cmd *cobra.Command, args []string) {
+	results = make(chan *result, putTotal)
+	requests := make(chan *etcdserverpb.PutRequest, putTotal)
+	bar = pb.New(putTotal)
+
+	k, v := mustRandBytes(keySize), mustRandBytes(valSize)
+
+	conns := make([]*grpc.ClientConn, totalConns)
+	for i := range conns {
+		conns[i] = mustCreateConn()
+	}
+
+	clients := make([]etcdserverpb.KVClient, totalClients)
+	for i := range clients {
+		clients[i] = etcdserverpb.NewKVClient(conns[i%int(totalConns)])
+	}
+
+	bar.Format("Bom !")
+	bar.Start()
+
+	for i := range clients {
+		wg.Add(1)
+		go doPut(clients[i], requests)
+	}
+
+	start := time.Now()
+	for i := 0; i < putTotal; i++ {
+		r := &etcdserverpb.PutRequest{
+			Key:   k,
+			Value: v,
+		}
+		requests <- r
+	}
+	close(requests)
+
+	wg.Wait()
+
+	bar.Finish()
+	printReport(putTotal, results, time.Now().Sub(start))
+}
+
+func doPut(client etcdserverpb.KVClient, requests <-chan *etcdserverpb.PutRequest) {
+	defer wg.Done()
+
+	for r := range requests {
+		st := time.Now()
+		_, err := client.Put(context.Background(), r)
+
+		var errStr string
+		if err != nil {
+			errStr = err.Error()
+		}
+		results <- &result{
+			errStr:   errStr,
+			duration: time.Now().Sub(st),
+		}
+		bar.Increment()
+	}
+}

tools/benchmark/cmd/range.go

@@ -0,0 +1,113 @@
+// Copyright 2015 CoreOS, Inc.
+//
+// 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 cmd
+
+import (
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/cheggaaa/pb"
+	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/spf13/cobra"
+	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
+	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"
+	"github.com/coreos/etcd/etcdserver/etcdserverpb"
+)
+
+// rangeCmd represents the range command
+var rangeCmd = &cobra.Command{
+	Use:   "range key [end-range]",
+	Short: "Benchmark range",
+
+	Run: rangeFunc,
+}
+
+var (
+	rangeTotal int
+)
+
+func init() {
+	RootCmd.AddCommand(rangeCmd)
+	rangeCmd.Flags().IntVar(&rangeTotal, "total", 10000, "Total number of range requests")
+}
+
+func rangeFunc(cmd *cobra.Command, args []string) {
+	if len(args) == 0 || len(args) > 2 {
+		fmt.Fprintln(os.Stderr, cmd.Usage())
+		os.Exit(1)
+	}
+
+	k := []byte(args[0])
+	var end []byte
+	if len(args) == 1 {
+		end = []byte(args[1])
+	}
+
+	results = make(chan *result, rangeTotal)
+	requests := make(chan *etcdserverpb.RangeRequest, rangeTotal)
+	bar = pb.New(rangeTotal)
+
+	conns := make([]*grpc.ClientConn, totalConns)
+	for i := range conns {
+		conns[i] = mustCreateConn()
+	}
+
+	clients := make([]etcdserverpb.KVClient, totalClients)
+	for i := range clients {
+		clients[i] = etcdserverpb.NewKVClient(conns[i%int(totalConns)])
+	}
+
+	bar.Format("Bom !")
+	bar.Start()
+
+	for i := range clients {
+		wg.Add(1)
+		go doRange(clients[i], requests)
+	}
+
+	start := time.Now()
+	for i := 0; i < rangeTotal; i++ {
+		r := &etcdserverpb.RangeRequest{
+			Key:      k,
+			RangeEnd: end,
+		}
+		requests <- r
+	}
+	close(requests)
+
+	wg.Wait()
+
+	bar.Finish()
+	printReport(rangeTotal, results, time.Now().Sub(start))
+}
+
+func doRange(client etcdserverpb.KVClient, requests <-chan *etcdserverpb.RangeRequest) {
+	defer wg.Done()
+
+	for req := range requests {
+		st := time.Now()
+		_, err := client.Range(context.Background(), req)
+
+		var errStr string
+		if err != nil {
+			errStr = err.Error()
+		}
+		results <- &result{
+			errStr:   errStr,
+			duration: time.Now().Sub(st),
+		}
+		bar.Increment()
+	}
+}

tools/benchmark/cmd/report.go

@@ -14,7 +14,7 @@
 
 // the file is borrowed from github.com/rakyll/boom/boomer/print.go
 
-package main
+package cmd
 
 import (
 	"fmt"

tools/benchmark/cmd/root.go

@@ -0,0 +1,57 @@
+// Copyright 2015 CoreOS, Inc.
+//
+// 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 cmd
+
+import (
+	"fmt"
+	"os"
+	"sync"
+
+	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/cheggaaa/pb"
+	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/spf13/cobra"
+)
+
+// This represents the base command when called without any subcommands
+var RootCmd = &cobra.Command{
+	Use:   "benchmark",
+	Short: "A low-level benchmark tool for etcd3",
+	Long: `benchmark is a low-level benchmakr tool for etcd3.
+It uses gRPC client directly and does not depend on 
+etcd client libray.
+	`,
+}
+
+var (
+	endpoints    string
+	totalConns   uint
+	totalClients uint
+
+	bar     *pb.ProgressBar
+	results chan *result
+	wg      sync.WaitGroup
+)
+
+func init() {
+	RootCmd.PersistentFlags().StringVar(&endpoints, "endpoint", "127.0.0.1:2378", "comma-separated gRPC endpoints")
+	RootCmd.PersistentFlags().UintVar(&totalConns, "conns", 1, "Total number of gRPC connections")
+	RootCmd.PersistentFlags().UintVar(&totalClients, "clients", 1, "Total number of gRPC clients")
+}
+
+func Execute() {
+	if err := RootCmd.Execute(); err != nil {
+		fmt.Println(err)
+		os.Exit(-1)
+	}
+}

tools/benchmark/cmd/util.go

@@ -0,0 +1,42 @@
+// Copyright 2015 CoreOS, Inc.
+//
+// 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 cmd
+
+import (
+	"crypto/rand"
+	"fmt"
+	"os"
+
+	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"
+)
+
+func mustCreateConn() *grpc.ClientConn {
+	conn, err := grpc.Dial(endpoints)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "dial error: %v\n", err)
+		os.Exit(1)
+	}
+	return conn
+}
+
+func mustRandBytes(n int) []byte {
+	rb := make([]byte, n)
+	_, err := rand.Read(rb)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "failed to generate value: %v\n", err)
+		os.Exit(1)
+	}
+	return rb
+}

tools/benchmark/main.go

@@ -0,0 +1,29 @@
+// Copyright 2015 CoreOS, Inc.
+//
+// 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 (
+	"fmt"
+	"os"
+
+	"github.com/coreos/etcd/tools/benchmark/cmd"
+)
+
+func main() {
+	if err := cmd.RootCmd.Execute(); err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(-1)
+	}
+}

tools/v3benchmark/get.go

@@ -1,57 +0,0 @@
-// Copyright 2015 CoreOS, Inc.
-//
-// 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 (
-	"time"
-
-	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
-	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"
-	"github.com/coreos/etcd/etcdserver/etcdserverpb"
-)
-
-func benchGet(conn *grpc.ClientConn, key, rangeEnd []byte, n, c int) {
-	wg.Add(c)
-	requests := make(chan struct{}, n)
-
-	for i := 0; i < c; i++ {
-		go get(etcdserverpb.NewKVClient(conn), key, rangeEnd, requests)
-	}
-
-	for i := 0; i < n; i++ {
-		requests <- struct{}{}
-	}
-	close(requests)
-}
-
-func get(client etcdserverpb.KVClient, key, end []byte, requests <-chan struct{}) {
-	defer wg.Done()
-	req := &etcdserverpb.RangeRequest{Key: key, RangeEnd: end}
-
-	for _ = range requests {
-		st := time.Now()
-		_, err := client.Range(context.Background(), req)
-
-		var errStr string
-		if err != nil {
-			errStr = err.Error()
-		}
-		results <- &result{
-			errStr:   errStr,
-			duration: time.Now().Sub(st),
-		}
-		bar.Increment()
-	}
-}

tools/v3benchmark/main.go

@@ -1,93 +0,0 @@
-// Copyright 2015 CoreOS, Inc.
-//
-// 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"
-	"fmt"
-	"os"
-	"strconv"
-	"sync"
-	"time"
-
-	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/cheggaaa/pb"
-	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"
-)
-
-var (
-	bar     *pb.ProgressBar
-	results chan *result
-	wg      sync.WaitGroup
-)
-
-func main() {
-	var (
-		c, n int
-		url  string
-		size int
-	)
-
-	flag.IntVar(&c, "c", 50, "number of connections")
-	flag.IntVar(&n, "n", 200, "number of requests")
-	flag.IntVar(&size, "s", 128, "size of put request")
-	// TODO: config the number of concurrency in each connection
-	flag.StringVar(&url, "u", "127.0.0.1:12379", "etcd server endpoint")
-	flag.Parse()
-	if flag.NArg() < 1 {
-		flag.Usage()
-		os.Exit(1)
-	}
-
-	var act string
-	if act = flag.Args()[0]; act != "get" && act != "put" {
-		fmt.Printf("unsupported action %v\n", act)
-		os.Exit(1)
-	}
-
-	conn, err := grpc.Dial(url)
-	if err != nil {
-		fmt.Errorf("dial error: %v", err)
-		os.Exit(1)
-	}
-
-	results = make(chan *result, n)
-	bar = pb.New(n)
-	bar.Format("Bom !")
-	bar.Start()
-
-	start := time.Now()
-
-	if act == "get" {
-		var rangeEnd []byte
-		key := []byte(flag.Args()[1])
-		if len(flag.Args()) > 2 {
-			rangeEnd = []byte(flag.Args()[2])
-		}
-		benchGet(conn, key, rangeEnd, n, c)
-	} else if act == "put" {
-		key := []byte(flag.Args()[1])
-		// number of different keys to put into etcd
-		kc, err := strconv.ParseInt(flag.Args()[2], 10, 32)
-		if err != nil {
-			panic(err)
-		}
-		benchPut(conn, key, int(kc), n, c, size)
-	}
-
-	wg.Wait()
-
-	bar.Finish()
-	printReport(n, results, time.Now().Sub(start))
-}

tools/v3benchmark/put.go

@@ -1,79 +0,0 @@
-// Copyright 2015 CoreOS, Inc.
-//
-// 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 (
-	"crypto/rand"
-	"encoding/binary"
-	"fmt"
-	"os"
-	"time"
-
-	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
-	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"
-	"github.com/coreos/etcd/etcdserver/etcdserverpb"
-)
-
-func benchPut(conn *grpc.ClientConn, key []byte, kc, n, c, size int) {
-	wg.Add(c)
-	requests := make(chan *etcdserverpb.PutRequest, n)
-
-	v := make([]byte, size)
-	_, err := rand.Read(v)
-	if err != nil {
-		fmt.Printf("failed to generate value: %v\n", err)
-		os.Exit(1)
-		return
-	}
-
-	for i := 0; i < c; i++ {
-		go put(etcdserverpb.NewKVClient(conn), requests)
-	}
-
-	suffixb := make([]byte, 8)
-	suffix := 0
-	for i := 0; i < n; i++ {
-		binary.BigEndian.PutUint64(suffixb, uint64(suffix))
-		r := &etcdserverpb.PutRequest{
-			Key:   append(key, suffixb...),
-			Value: v,
-		}
-		requests <- r
-		if suffix > kc {
-			suffix = 0
-		}
-		suffix++
-	}
-	close(requests)
-}
-
-func put(client etcdserverpb.KVClient, requests <-chan *etcdserverpb.PutRequest) {
-	defer wg.Done()
-
-	for r := range requests {
-		st := time.Now()
-		_, err := client.Put(context.Background(), r)
-
-		var errStr string
-		if err != nil {
-			errStr = err.Error()
-		}
-		results <- &result{
-			errStr:   errStr,
-			duration: time.Now().Sub(st),
-		}
-		bar.Increment()
-	}
-}