Merge pull request #4 from tombokombo/master

rework

Author: yaron2

Dockerfile

@@ -1,12 +1,13 @@
-FROM golang:1.9  AS builder
+FROM golang:1.10-alpine  AS builder
 
 RUN mkdir /app
-ADD . /app/
+ADD ./main.go /app/
 WORKDIR /app
-RUN go get github.com/go-sql-driver/mysql
+RUN apk update && apk add git && go get github.com/go-sql-driver/mysql && go get github.com/namsral/flag
 RUN go build -o main .
 
 FROM alpine:3.10
-
-COPY --from=builder /app/main /usr/bin/mysql-load-generator
-CMD ["/usr/bin/mysql-load-generator"]
+RUN mkdir /app
+COPY --from=builder /app/main /app/mysql-loader
+COPY config.json /app/config.json
+CMD ["/app/mysql-loader"]

Readme.md

@@ -1,6 +1,7 @@
 # MemSQL / MySQL Load Tester
 
 This is a tool for generating load tests based on custom queries for MemSQL / MySQL databases.
+This fork brings substitution, query timeouts, wait group timeout, connection pool, periodic qps and latency logs, higher performance
 
 ## How does it work?
 
@@ -11,13 +12,13 @@ Json configuration guide:
 
 * connectionString (string) - the connection string to the MySQL / MemSQL Database
 * requestsPerSecond (int) - how many statements per second to execute
-* printLogs (bool) - print queries executed including time taken for each query and a total average
+* printLogs (bool) - print queries executed 
 * timeToRun (int) - how long to run the test in seconds
-* queries (string array) - the queries to run
+* queries (string array) - the queries to run, can use substitution placeholder
+* substitution ( map ) - specify placeholder for substitution with random number according specified min,max (int) rules
 
 ## Run as Docker
-
-Simply build using the Dockerfile.
-Load tests are a great fit for [Azure Container Instances].
-
-[Azure Container Instances]: https://azure.microsoft.com/en-us/services/container-instances/
+Prepare config.json and run
+```
+docker run --network=host -v $(pwd)/config.json:/app/config.json -it --rm tombokombo/mysql-loader:latest
+```

config.json

@@ -1,9 +1,13 @@
 {
     "connectionString": "USER:PASSWORD@tcp(HOST:PORT)/DB",
-    "requestsPerSecond": 1,
+    "requestsPerSecond": 4500,
     "printLogs": true,
-    "timeToRun": 3,
+    "timeToRunInSeconds": 1200,
+    "poolConnections": 1000,
+    "dryRun"  : false,
+    "queryTimeout": 120,
     "queries" : [
-        "SELECT * from Table"
-    ]
+    "SELECT * FROM test LIMIT WHERE id=:rand: LIMIT 1"
+    ],
+    "substitution" : [ {"key": ":rand:", "min": 1 ,"max": 3182356 }, {"key": ":rand_tiny:", "min": 1 ,"max": 10 } ]
 }

main.go

@@ -1,42 +1,105 @@
 package main
 
 import (
+	"context"
 	"database/sql"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
 	"math/rand"
+	"os"
 	"strconv"
+	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
+	"github.com/namsral/flag"
+
 	_ "github.com/go-sql-driver/mysql"
 )
 
+type Substitution struct {
+	Key string `json:"key"`
+	Min int    `json:"min"`
+	Max int    `json:"max"`
+}
+
 type Config struct {
-	RequestsPerSecond  int      `json:"requestsPerSecond"`
-	Queries            []string `json:"queries"`
-	ConnectionString   string   `json:"connectionString"`
-	PrintLogs          bool     `json:"printLogs"`
-	TimeToRunInSeconds int      `json:"timeToRunInSeconds"`
+	RequestsPerSecond  int            `json:"requestsPerSecond"`
+	Queries            []string       `json:"queries"`
+	ConnectionString   string         `json:"connectionString"`
+	PrintLogs          bool           `json:"printLogs"`
+	TimeToRunInSeconds int            `json:"timeToRunInSeconds"`
+	PoolConnections    int            `json:"poolConnections"`
+	DryRun             bool           `json:"dryRun"`
+	QueryTimeout       int            `json:"queryTimeout"`
+	ConnectionLifetime time.Duration  `json:"connectionLifeTime"`
+	Substitution       []Substitution `json:"substitution"`
 }
 
-var config Config
+var config = Config{
+	RequestsPerSecond:  100,
+	Queries:            make([]string, 0),
+	ConnectionString:   "",
+	PrintLogs:          true,
+	TimeToRunInSeconds: 10,
+	PoolConnections:    10,
+	DryRun:             true,
+	QueryTimeout:       10,
+	ConnectionLifetime: 0,
+	Substitution:       []Substitution{},
+}
+var configFilePath string
 var db *sql.DB
 var queryTimesInMS []int
+var ops int64 = 0
+var latency int64 = 0
 
 var wg sync.WaitGroup
 
 func main() {
-	loadConfig()
+	flag.StringVar(&configFilePath, "configFile", "/app/config.json", "path to config file")
+	flag.Parse()
+	go printQps()
+	loadConfig(configFilePath)
 	openMemSQLConnection()
 	dispatchQueries()
-	displayAverageQueryTime()
+}
+
+func printQps() {
+	ticker := time.NewTicker(5 * time.Second)
+	for range ticker.C {
+		loadedOps := atomic.LoadInt64(&ops)
+		atomic.StoreInt64(&ops, 0)
+		loadedLatency := atomic.LoadInt64(&latency)
+		atomic.StoreInt64(&latency, 0)
+		fmt.Println("qps: ", loadedOps/5, ", average latency: ", time.Duration(loadedLatency/loadedOps))
+
+	}
+
+}
+
+func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
+	c := make(chan struct{})
+	go func() {
+		defer close(c)
+		wg.Wait()
+	}()
+	select {
+	case <-c:
+		return false // completed normally
+	case <-time.After(timeout):
+		return true // timed out
+	}
 }
 
 func openMemSQLConnection() {
 	var err error
 	db, err = sql.Open("mysql", config.ConnectionString)
+	db.SetConnMaxLifetime(config.ConnectionLifetime)
+	db.SetMaxIdleConns(config.PoolConnections)
+	db.SetMaxOpenConns(config.PoolConnections)
 
 	if err != nil {
 		panic(err.Error())
@@ -48,45 +111,30 @@ func openMemSQLConnection() {
 		panic(err.Error())
 	}
 
-	fmt.Println("Connection succeeded")
-}
-
-func trackFuncTime(start time.Time) {
-	elapsed := time.Since(start)
-	queryTimesInMS = append(queryTimesInMS, int(elapsed/time.Millisecond))
-
-	if config.PrintLogs {
-		fmt.Println(elapsed)
-	}
+	fmt.Println("Connection succeeded with pool size", config.PoolConnections)
 }
 
 func executeQuery(query string) {
 	defer wg.Done()
-	defer trackFuncTime(time.Now())
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(config.QueryTimeout)*time.Second)
+	defer cancel()
 
-	_, err := db.Exec(query)
-
-	if err != nil {
-		fmt.Println(err.Error())
+	for _, subs := range config.Substitution {
+		query = strings.Replace(query, subs.Key, strconv.Itoa(rand.Intn(subs.Max-subs.Min)+subs.Min), 1)
 	}
 
 	if config.PrintLogs {
-		fmt.Println("Executed " + query)
+		fmt.Println("Executed ", query)
 	}
-}
-
-func displayAverageQueryTime() {
-	if config.PrintLogs {
-		var totalvalues int
-
-		for _, val := range queryTimesInMS {
-			totalvalues += val
+	if config.DryRun == false {
+		start := time.Now()
+		_, err := db.ExecContext(ctx, query)
+		if err == nil {
+			atomic.AddInt64(&ops, 1)
+			atomic.AddInt64(&latency, int64(time.Since(start)))
+		} else {
+			fmt.Fprintf(os.Stderr, "query err: %s\n", err.Error())
 		}
-
-		var count = len(queryTimesInMS)
-		sum := totalvalues / count
-
-		fmt.Println("Average time: " + strconv.Itoa(sum) + "ms")
 	}
 }
 
@@ -97,23 +145,44 @@ func dispatchQueries() {
 	shouldEnd := start.Add(time.Second * time.Duration(config.TimeToRunInSeconds))
 
 	for time.Now().Before(shouldEnd) {
-		for i := 0; i < config.RequestsPerSecond; i++ {
-			wg.Add(1)
-			go executeQuery(config.Queries[rand.Intn(len(config.Queries))])
+		l := len(config.Queries)
+		for i := 0; i < config.RequestsPerSecond/l; i++ {
+			wg.Add(l)
+			for q := 0; q < l; q++ {
+				go executeQuery(config.Queries[q])
+			}
 		}
 
 		time.Sleep(1 * time.Second)
 	}
-
-	wg.Wait()
+	fmt.Println("wait")
+	if waitTimeout(&wg, 10*time.Second) {
+		fmt.Println("Timed out waiting for wait group")
+	} else {
+		fmt.Println("Wait group finished")
+	}
+	fmt.Println("end wait")
+	defer db.Close()
 }
 
-func loadConfig() {
-	data, err := ioutil.ReadFile("config.json")
+func loadConfig(configFilePath string) {
+	data, err := ioutil.ReadFile(configFilePath)
+	if err != nil {
+		panic(err.Error() + " use -configFile flag")
+	}
+
+	err = json.Unmarshal(data, &config)
 	if err != nil {
 		panic(err)
 	}
+	fmt.Println("LOADED CONFIG")
+	pretty, err := json.MarshalIndent(config, "", "    ")
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("%s\n", string(pretty))
 
-	json.Unmarshal(data, &config)
-	fmt.Println("Config loaded")
+	if len(config.Queries) < 1 {
+		panic("no query specified")
+	}
 }