package main

import (
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"os"
	"strings"
	"time"
	"math/rand"
	"sync"
	"strconv"
)

type (
	subMap map[int64]int64
)

const (
	CHARSET = "utf8mb4"
)

var (
	db         *sql.DB
	dbErr      error
	errCounter int
	// params
	help         bool
	output       bool
	host         string
	port         int
	databaseName string
	username     string
	passwd       string
	concurrency  int
	runTime      int
	sbtable_num  int
	sbtable_size int
	dbTimeout    int
)

func usage() {
	fmt.Fprintf(os.Stderr, `sd-benchmark version: 1.0.1
Desc : MySQL bench mark tools 
Usage: sd-benchmark [-h Connect to host.] [-P Port number to use for connection] [-d database name] [-u User for login if not current user.] [-p Password to use when connecting to server] [-r running time] [-f SQL file] [-t timeout] [-o true | false]

Options:
`)
	flag.PrintDefaults()
}

var tableNameMap = struct{
    sync.Mutex
    m map[int]string
}{m: make(map[int]string)}

func init() {
	//version = flag.Bool("V", false, "show version and exit")
	flag.BoolVar(&help, "H", false, "this help")
	flag.StringVar(&host, "h", "127.0.0.1", "Connect to host.")
	flag.IntVar(&port, "P", 3306, "Port number to use for connection.")
	flag.StringVar(&databaseName, "d", "information_schema", "Database to use.")
	flag.StringVar(&username, "u", "root", "User for login if not current user.")
	flag.StringVar(&passwd, "p", "", "Password to use when connecting to server. (default \"\")")
	flag.IntVar(&concurrency, "c", 64, "Number of threads.")
	flag.IntVar(&runTime, "r", 600, "running time.")
	flag.IntVar(&sbtable_num, "f", 10, "Number of sysbench tables. (default 10")
	flag.IntVar(&sbtable_size, "g", 100, "Number of keys in sysbench tables. (default 100")
	flag.IntVar(&dbTimeout, "t", 60, "connect MySQL timeout.")
	flag.BoolVar(&output, "o", false, "output error info.")
	flag.Usage = usage
}

func RandInt(rd *rand.Rand, start int, end int ) int {
	return rd.Intn(end - start + 1) + start
}

func main() {

	flag.Parse()
	if help {
		flag.Usage()
		return
	}

	// initialize db
	getDB(host, port, databaseName, username, passwd)
	defer db.Close()

	queueMap := make(map[int]subMap)
	for i := 1; i <= concurrency; i++ {
		queueMap[i] = make(subMap)
	}
	startTime := time.Now()
	fmt.Printf("\nINFO: %s start performance testing (%d threads) ... \n", startTime.Format("2006-01-02 15:04:05"), concurrency)
	for i := 1; i <= concurrency; i++ {
		go subProcess(int64(i), queueMap[i])
	}

	select {
	case <-time.After(time.Duration(runTime) * time.Second):
		fmt.Printf("INFO: %s DONE.\n", time.Now().Format("2006-01-02 15:04:05"))
		endTime := time.Now()
		elapsed := (endTime.UnixNano() - startTime.UnixNano()) / int64(time.Millisecond)
		analyzePerformanceResult(queueMap, float64(elapsed)/1000)
	}
}

func getDB(host string, port int, dbName string, user string, passwd string) {
	dbDSN := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&timeout=%ds&readTimeout=%ds", user, passwd, host, port, dbName, CHARSET, dbTimeout, dbTimeout)

	db, dbErr = sql.Open("mysql", dbDSN)
	//defer MysqlDb.Close();
	if dbErr != nil {
		if output {
			fmt.Println("dbDSN: " + dbDSN)
			panic("ERROR: Data source configuration is incorrect, err: " + dbErr.Error())
		}
	}

	db.SetMaxOpenConns(concurrency * 2)
	db.SetMaxIdleConns(100)
	db.SetConnMaxLifetime(600 * time.Second)

	if dbErr = db.Ping(); nil != dbErr {
		if output {
			fmt.Println("db conn=", dbDSN)
			panic("connect database failed ,err: " + dbErr.Error())
		}
	}
}

func getCurrentTableName(tableId int) string {
	tableNameMap.Lock();
	defer tableNameMap.Unlock();
	if tableName, ok := tableNameMap.m[tableId]; ok {
		return tableName
	} else {
		tableNameMap.m[tableId] = "sbtest" + strconv.Itoa(tableId)
		return  "sbtest" + strconv.Itoa(tableId)
	}
}

/* return previous name and current name */
func alterCurrentTableName(tableId int) (string, string) {
	tableNameMap.Lock();
	defer tableNameMap.Unlock();
	if oldTableName, ok := tableNameMap.m[tableId]; ok {
		newTableName := ""
		if strings.HasSuffix(oldTableName, "tmp") {
			newTableName = strings.TrimSuffix(oldTableName, "tmp")			
		} else {
			newTableName = oldTableName + "tmp"
		}
		tableNameMap.m[tableId] = newTableName
		return oldTableName, newTableName
	} else {
		tableNameMap.m[tableId] = "sbtest" + strconv.Itoa(tableId)
		return  "", "sbtest" + strconv.Itoa(tableId)
	}
}

func subProcess(id int64, costMap subMap) {
	seed := rand.NewSource(id)
	rd := rand.New(seed)
	for {
		dice := rd.Intn(1000)
		sql := ""
		tableId := RandInt(rd, 1, sbtable_num)
		if dice == 0 {
			oldTableName, newTableName := alterCurrentTableName(tableId)
			if len(oldTableName) == 0 {
				sql = "select 1"
			} else {
				sql = "ALTER TABLE " + oldTableName + " RENAME " + newTableName
			}
		} else {
			keyId := RandInt(rd, 1, sbtable_size)
			tableName := getCurrentTableName(tableId)
			sql = "UPDATE " + tableName + " SET k=k+1 WHERE id=" + strconv.Itoa(keyId)
		}
		start := time.Now().UnixNano()
		_, err := db.Exec(sql)
		end := time.Now().UnixNano()
		if err != nil {
			if output {
				fmt.Printf("ERROR: query  exception sql: %s, err: %s\n", sql, err)
			}
			errCounter++
			continue
		}
		costTime := (end - start) / int64(time.Millisecond)
		if _, ok := costMap[costTime]; ok {
			costMap[costTime]++
		} else {
			costMap[costTime] = 1
		}
	}
}

func analyzePerformanceResult(queueMap map[int]subMap, elapsed float64) {

	fmt.Printf("INFO: %s start analysis report ... \n", time.Now().Format("2006-01-02 15:04:05"))
	var QPS, costTime, minTime, maxTime int64
	minTime = 20
	for _, item := range queueMap {
		//fmt.Printf("\nqueue %d:\n", queue)
		var tmpQPS, tmpCostTime, tmpMinTime, tmpMaxTime int64
		tmpMinTime = 20
		for k, v := range item {
			tmpQPS += v
			tmpCostTime += k * v
			if k > tmpMaxTime {
				tmpMaxTime = k
			}
			if k < tmpMinTime {
				tmpMinTime = k
			}
		}
		QPS += tmpQPS
		costTime += tmpCostTime
		if tmpMaxTime > maxTime {
			maxTime = tmpMaxTime
		}
		if tmpMinTime < minTime {
			minTime = tmpMinTime
		}
		//fmt.Println("qps=", tmpQPS)
		//fmt.Println("costTime=", tmpCostTime)
		//fmt.Println("minResponseTime=", tmpMinTime)
		//fmt.Println("maxResponseTime=", tmpMaxTime)
	}
	fmt.Printf("INFO: %s DONE.\n", time.Now().Format("2006-01-02 15:04:05"))

	fmt.Printf(`
SQL statistics:
    queries performed:
         read :                          - (nonsupport)
         write:                          - (nonsupport)
         other:                          - (nonsupport)
         total:                          %d
       queries:                          %d (%f per sec.)
ignored errors:                          %d (%f per sec.)


General statistics:
    number of threads:                   %d
    conn db timeout:                     %ds
    read db timeout:                     %ds
    total time:                          %fs


Latency (ms):
         min:                            %d
         avg:                            %f
         max:                            %d
         sum:                            %d
`, QPS+int64(errCounter), QPS, float64(QPS)/elapsed, errCounter, float64(errCounter)/elapsed, concurrency, dbTimeout, dbTimeout, elapsed, minTime, float32(costTime)/float32(QPS), maxTime, costTime)
}