redisLoadTest/redisLoadTest.go

package main

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

var ctx = context.Background()

/**
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 := 0; i < size; i++ {
		key := i
		value := "ThisIsATestStringThatShouldBeReplacedWithSomethingMoreRandomMaybeFromAnInputOrCSV"
		go setRedis(rdbPtr, key, value, &waitGroup)
		waitGroup.Add(1)
	}
	waitGroup.Wait()
	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) {
	x := int(math.Abs(60 / float64(rate) * 1000))
	rateLimiter := time.Tick(time.Millisecond * time.Duration(x))

	fmt.Println("Starting Test at", rateLimiter, "milliseconds per request...")
	//todo add end to test
	for {
		<-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) {
	burstRateLimiter := time.Tick(time.Second)
	//todo add end to test
	fmt.Println("Starting burst test at", rate, "requests per second...")
	for {
		<-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
	_, err := rdb.Do("Set", key, value)
	if err != nil {
		defer waitGroup.Done()
		panic(err)
	}
	defer waitGroup.Done()
}

/**
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) *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())
			}
			return c, err
		},
	}
}

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", 50, "Test rate limit, default 50/sec")
	burstPtr := flag.Bool("burst", true, "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")

	host := *hostPtr
	//todo add username/password support
	//username := *usernamePtr
	//password := *passwordPtr
	//db := *dbPtr
	initializeDB := *initializeDBPtr
	rate := *ratePtr
	burst := *burstPtr
	keyMaxValue := *maxEntriesPtr

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

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