From de899fbd826bc594a680bf541674601d715beb3b Mon Sep 17 00:00:00 2001
From: "Sahdev P. Zala" <spzala@us.ibm.com>
Date: Sun, 21 Jan 2018 23:43:59 -0500
Subject: [PATCH] etcdctl/check: create new check command for memory usage

Create a new command similar to check perf that can check the memory
consumption for putting different workloads. Return user with a message
that whether there are enough memory for a given workload with pass
or fail.

Fixed #9121
---
 etcdctl/ctlv3/command/check.go | 158 ++++++++++++++++++++++++++++++++-
 1 file changed, 156 insertions(+), 2 deletions(-)

diff --git a/etcdctl/ctlv3/command/check.go b/etcdctl/ctlv3/command/check.go
index 2ad568217a07..2b7c6b81f271 100644
--- a/etcdctl/ctlv3/command/check.go
+++ b/etcdctl/ctlv3/command/check.go
@@ -21,6 +21,7 @@ import (
 	"math"
 	"math/rand"
 	"os"
+	"runtime"
 	"sync"
 	"time"
 
@@ -33,8 +34,10 @@ import (
 )
 
 var (
-	checkPerfLoad   string
-	checkPerfPrefix string
+	checkPerfLoad        string
+	checkPerfPrefix      string
+	checkDatascaleLoad   string
+	checkDatascalePrefix string
 )
 
 type checkPerfCfg struct {
@@ -67,6 +70,35 @@ var checkPerfCfgMap = map[string]checkPerfCfg{
 	},
 }
 
+type checkDatascaleCfg struct {
+	limit   int
+	kvSize  int
+	clients int
+}
+
+var checkDatascaleCfgMap = map[string]checkDatascaleCfg{
+	"s": {
+		limit:   10000,
+		kvSize:  1024,
+		clients: 50,
+	},
+	"m": {
+		limit:   100000,
+		kvSize:  1024,
+		clients: 200,
+	},
+	"l": {
+		limit:   1000000,
+		kvSize:  1024,
+		clients: 500,
+	},
+	"xl": {
+		limit:   30000000,
+		kvSize:  1024,
+		clients: 1000,
+	},
+}
+
 // NewCheckCommand returns the cobra command for "check".
 func NewCheckCommand() *cobra.Command {
 	cc := &cobra.Command{
@@ -75,6 +107,7 @@ func NewCheckCommand() *cobra.Command {
 	}
 
 	cc.AddCommand(NewCheckPerfCommand())
+	cc.AddCommand(NewCheckDatascaleCommand())
 
 	return cc
 }
@@ -216,3 +249,124 @@ func newCheckPerfCommand(cmd *cobra.Command, args []string) {
 		os.Exit(ExitError)
 	}
 }
+
+// NewCheckDatascaleCommand returns the cobra command for "check datascale".
+func NewCheckDatascaleCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "datascale",
+		Short: "Check the availability of enough memory for holding data",
+		Run:   newCheckDatascaleCommand,
+	}
+
+	cmd.Flags().StringVar(&checkDatascaleLoad, "load", "s", "The datascale check's workload model. Accepted workloads: s(small), m(medium), l(large), xl(xLarge)")
+	cmd.Flags().StringVar(&checkDatascalePrefix, "prefix", "/etcdctl-check-datascale/", "The prefix for writing the datascale check's keys.")
+
+	return cmd
+}
+
+// newCheckDatascaleCommand executes the "check datascale" command.
+func newCheckDatascaleCommand(cmd *cobra.Command, args []string) {
+	var checkDatascaleAlias = map[string]string{
+		"s": "s", "small": "s",
+		"m": "m", "medium": "m",
+		"l": "l", "large": "l",
+		"xl": "xl", "xLarge": "xl",
+	}
+
+	model, ok := checkDatascaleAlias[checkDatascaleLoad]
+	if !ok {
+		ExitWithError(ExitBadFeature, fmt.Errorf("unknown load option %v", checkDatascaleLoad))
+	}
+	cfg := checkDatascaleCfgMap[model]
+
+	requests := make(chan v3.Op, cfg.clients)
+
+	cc := clientConfigFromCmd(cmd)
+	clients := make([]*v3.Client, cfg.clients)
+	for i := 0; i < cfg.clients; i++ {
+		clients[i] = cc.mustClient()
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	resp, err := clients[0].Get(ctx, checkDatascalePrefix, v3.WithPrefix(), v3.WithLimit(1))
+	cancel()
+	if err != nil {
+		ExitWithError(ExitError, err)
+	}
+	if len(resp.Kvs) > 0 {
+		ExitWithError(ExitInvalidInput, fmt.Errorf("prefix %q has keys. Delete with etcdctl del --prefix %s first.", checkDatascalePrefix, checkDatascalePrefix))
+	}
+
+	ksize, vsize := cfg.kvSize, cfg.kvSize
+	k, v := make([]byte, ksize), string(make([]byte, vsize))
+
+	r := report.NewReport("%4.4f")
+	var wg sync.WaitGroup
+
+	var statsBefore runtime.MemStats
+	runtime.ReadMemStats(&statsBefore)
+	totalOSMemory := statsBefore.Sys
+	totalAllocBefore := statsBefore.Alloc
+	wg.Add(len(clients))
+	for i := range clients {
+		go func(c *v3.Client) {
+			defer wg.Done()
+			for op := range requests {
+				st := time.Now()
+				_, derr := c.Do(context.Background(), op)
+				r.Results() <- report.Result{Err: derr, Start: st, End: time.Now()}
+			}
+		}(clients[i])
+	}
+
+	go func() {
+		for i := 0; i < cfg.limit; i++ {
+			binary.PutVarint(k, int64(rand.Int63n(math.MaxInt64)))
+			requests <- v3.OpPut(checkDatascalePrefix+string(k), v)
+		}
+		close(requests)
+	}()
+
+	sc := r.Stats()
+	wg.Wait()
+	close(r.Results())
+	s := <-sc
+
+	// check memory consumption
+	var statsAfter runtime.MemStats
+	runtime.ReadMemStats(&statsAfter)
+	totalAllocAfter := statsAfter.Alloc
+	diffAlloc := totalAllocAfter - totalAllocBefore
+	// rounded percent of memory consumed
+	percentAlloc := (diffAlloc*100) / totalOSMemory
+
+	ctx, cancel = context.WithCancel(context.Background())
+	_, err = clients[0].Delete(ctx, checkDatascalePrefix, v3.WithPrefix())
+	cancel()
+	if err != nil {
+		ExitWithError(ExitError, err)
+	}
+
+	ok = true
+	if len(s.ErrorDist) != 0 {
+		fmt.Println("FAIL: too many errors")
+		for k, v := range s.ErrorDist {
+			fmt.Printf("FAIL: ERROR(%v) -> %d\n", k, v)
+		}
+		ok = false
+	}
+
+	if percentAlloc > 60 { // leaves less than 40 percent of memory
+		fmt.Println("FAIL: Memory usage is too high. Percent memory used is :", percentAlloc)
+		ok = false
+	} else {
+		fmt.Println("PASS: Memory usage is optimal. Percent memory used is : ", percentAlloc)
+	}
+
+	if ok {
+		fmt.Println("PASS")
+	} else {
+		fmt.Println("FAIL")
+		os.Exit(ExitError)
+	}
+}