redisLoadTest/redisLoadTest.go

package main

/*
 * ----------------------------------------------------------------------------
 * "THE BEER-WARE LICENSE" (Revision 42):
 * <mitch@nerdfortress.dev> wrote this file. As long as you retain this notice you
 * can do whatever you want with this stuff. If we meet some day, and you think
 * this stuff is worth it, you can buy me a beer in return
 * ----------------------------------------------------------------------------
 */

import (
	"flag"
	"fmt"
	"github.com/gomodule/redigo/redis"
	"github.com/montanaflynn/stats"
	"math"
	"math/rand"
	"sync"
	"time"
)

/**
buildTestData iterates through a goroutine of setRedis to build test data into the redis db
@param redisPtr connection to redis
@param size int number of entries into the database
*/
func buildTestData(rdbPtr *redis.Conn, size int) {
	fmt.Println("Initializing Test Data...")
	var waitGroup sync.WaitGroup
	for i := 1; i < 50000; i++ {
		key := i
		value := "ThisIsATestStringThatShouldBeReplacedWithSomethingMoreRandomMaybeFromAnInputOrCSV"
		setRedis(rdbPtr, key, value, &waitGroup)
	}
	fmt.Println("Database Initialized...")
}

/**
getRateTest calculates closest value for requests per milliseconds from inputted requests per second
Then starts goroutines for getRedis at that interval
@param rdbPtr connection is a pointer to the redis connection
@param rate int requests per second for the test
@param keyMaxValue int max value to query the database and expect a return
*/
func getRateTest(rdbPtr *redis.Conn, rate int, keyMaxValue int, duration time.Duration) {
	durationChannel := make(chan bool, 1)
	x := int(math.Abs(60 / float64(rate) * 10000000))
	rateLimiter := time.Tick(time.Duration(x))

	fmt.Println("Starting Test at", rateLimiter, "milliseconds per request...")
	go func() {
		for i := 0; i < int(duration.Seconds()); i++ {
			<-time.Tick(time.Second)
		}
		durationChannel <- true

	}()
	end := false
	for {
		if end {
			break
		}
		select {
		case <-durationChannel:
			fmt.Println("Closing Test...")
			end = true
		case <-rateLimiter:
			key := rand.Intn(keyMaxValue)
			go getRedis(rdbPtr, key)
		}
	}
	fmt.Println("Rate test concluded...")
}

/**
getBurstTest starts goroutines for getRedis at the burstRateLimit per second
@param rdbPtr connection is a pointer to the redis connection
@param rate int requests per second for the test
@param keyMaxValue int max value to query the database and expect a return
*/
func getBurstTest(rdbPtr *redis.Conn, rate int, keyMaxValue int, duration time.Duration) {
	durationChannel := make(chan bool, 1)
	burstRateLimiter := time.Tick(time.Second)
	go func() {
		for i := 0; i < int(duration.Seconds()); i++ {
			<-time.Tick(time.Second)
		}
		durationChannel <- true
	}()
	fmt.Println("Starting burst test at", rate, "requests per second...")
	end := false
	for {
		if end {
			break
		}
		select {
		case <-durationChannel:
			end = true
		case <-burstRateLimiter:
			for i := 0; i < rate; i++ {
				key := rand.Intn(keyMaxValue)
				go getRedis(rdbPtr, key)
			}
		}
	}
	fmt.Println("Burst test concluded...")
}

/**
setRateTest sets key/value into database
@param rdbPtr connection is a pointer to the redis connection
@param key int redis key for the storage
@param value string string for the value to store
*/
func setRedis(rdbPtr *redis.Conn, key int, value string, waitGroup *sync.WaitGroup) {
	rdb := *rdbPtr
	//defer waitGroup.Done()
	_, err := rdb.Do("Set", key, value)
	if err != nil {
		fmt.Println(err, key, value)
		panic(err)
	}
}

/**
getRedis queries the redis database for value of key
@param rdbPtr connection is a pointer to the redis connection
@param key int is key value to query in database
*/
func getRedis(rdbPtr *redis.Conn, key int) {
	rdb := *rdbPtr
	_, err := rdb.Do("Get", key)
	if err != nil {
		fmt.Println("Unhandled error:", err)
	}
}

/**
newPool builds redis database connection
@param host string host:port of redis server
returns redis.Connection of connection to database
returns err from database connection
*/
func newPool(host string, username string, password string) *redis.Pool {
	return &redis.Pool{
		MaxIdle:   80,
		MaxActive: 12000,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", host)
			if err != nil {
				panic(err.Error())
			}
			if username != "" && password != "" {
				_, err = c.Do("AUTH", password)
				if err != nil {
					c.Close()
					panic(err.Error())
				}
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}

func keepTotal(timesChan chan float64, errorChan chan bool, endChan chan bool) ([]float64, int) {
	var timesSlice []float64
	var errorRate int
	for {
		select {
		case <-endChan:
			break
		case timeValue := <-timesChan:
			timesSlice = append(timesSlice, timeValue)
		case <-errorChan:
			errorRate++
		}
	}
	return timesSlice, errorRate

}

/**
getResponseTimes takes a list of ints and returns mean, stddev, 95p, 99p
Uses stats library to calculate mean, stddev, 95p, 99p
*/
func getResponseTimes(times []float64) (float64, float64, float64, float64) {
	mean, _ := stats.Mean(times)
	stddev, _ := stats.StandardDeviation(times)
	perc99, _ := stats.Percentile(times, 99.00)
	perc95, _ := stats.Percentile(times, 95.00)
	return mean, stddev, perc99, perc95
}

func main() {
	hostPtr := flag.String("host", "", "Redis Server FQDN:port")
	usernamePtr := flag.String("username", "", "Redis Server username")
	passwordPtr := flag.String("password", "", "Redis user password")
	//dbPtr := flag.Int("db", 0, "Redis db, default 0")
	initializeDBPtr := flag.Bool("initialize", false, "Boolean initialize db, default false")
	ratePtr := flag.Int("rate", 5, "Test rate limit, default 50/sec")
	burstPtr := flag.Bool("burst", false, "Boolean burst test default true")
	//todo this should be off that same csv rather than an input
	maxEntriesPtr := flag.Int("dbEntries", 50000, "Test rate limit, default 50/sec")
	testDurationPtr := flag.Int("duration", 10, "Duration of each test")

	flag.Parse()
	host := *hostPtr
	username := *usernamePtr
	password := *passwordPtr
	//db := *dbPtr
	initializeDB := *initializeDBPtr
	rate := *ratePtr
	burst := *burstPtr
	keyMaxValue := *maxEntriesPtr
	durationInt := *testDurationPtr
	duration := time.Second * time.Duration(durationInt)

	pool := newPool(host, username, password)
	client := pool.Get()
	defer client.Close()

	if initializeDB {
		buildTestData(&client, 50000)
	}
	getRateTest(&client, rate, keyMaxValue, duration)
	if burst {
		getBurstTest(&client, rate, keyMaxValue, duration)
	}
	fmt.Println("Tests completed...")
}