From ee8c466b7873f3c5ff63458444160ebef5cef8ba Mon Sep 17 00:00:00 2001
From: mitch <mitch@nerdfortress.dev>
Date: Tue, 4 Jan 2022 19:19:31 -0500
Subject: [PATCH] Added auth support Improved pool handling Improved ticker for
 rate test

---
 redisLoadTest.go | 42 +++++++++++++++++++++++-------------------
 1 file changed, 23 insertions(+), 19 deletions(-)

diff --git a/redisLoadTest.go b/redisLoadTest.go
index 6345652..f72c830 100644
--- a/redisLoadTest.go
+++ b/redisLoadTest.go
@@ -32,9 +32,7 @@ func buildTestData(rdbPtr *redis.Conn, size int) {
 		key := i
 		value := "ThisIsATestStringThatShouldBeReplacedWithSomethingMoreRandomMaybeFromAnInputOrCSV"
 		setRedis(rdbPtr, key, value, &waitGroup)
-		//waitGroup.Add(1)
 	}
-	//waitGroup.Wait()
 	fmt.Println("Database Initialized...")
 }
 
@@ -47,19 +45,16 @@ Then starts goroutines for getRedis at that interval
 */
 func getRateTest(rdbPtr *redis.Conn, rate int, keyMaxValue int, duration time.Duration) {
 	durationChannel := make(chan bool, 1)
-	x := int(math.Abs(60 / float64(rate) * 1000))
-	rateLimiter := time.Tick(time.Millisecond * time.Duration(x))
+	x := int(math.Abs(60 / float64(rate) * 10000000))
+	rateLimiter := time.Tick(time.Duration(x))
 
 	fmt.Println("Starting Test at", rateLimiter, "milliseconds per request...")
-	//	select {
-	//	case <- time.After(10 * time.Second):
-	//		fmt.Fprintf(w, "hello\n")
 	go func() {
 		for i := 0; i < int(duration.Seconds()); i++ {
-			fmt.Println(i, int(duration.Seconds()))
 			<-time.Tick(time.Second)
 		}
 		durationChannel <- true
+
 	}()
 	end := false
 	for {
@@ -72,7 +67,6 @@ func getRateTest(rdbPtr *redis.Conn, rate int, keyMaxValue int, duration time.Du
 			end = true
 		case <-rateLimiter:
 			key := rand.Intn(keyMaxValue)
-			fmt.Println(key)
 			go getRedis(rdbPtr, key)
 		}
 	}
@@ -148,7 +142,7 @@ newPool builds redis database connection
 returns redis.Connection of connection to database
 returns err from database connection
 */
-func newPool(host string) *redis.Pool {
+func newPool(host string, username string, password string) *redis.Pool {
 	return &redis.Pool{
 		MaxIdle:   80,
 		MaxActive: 12000,
@@ -157,8 +151,19 @@ func newPool(host string) *redis.Pool {
 			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
+		},
 	}
 }
 
@@ -193,21 +198,20 @@ func getResponseTimes(times []float64) (float64, float64, float64, float64) {
 
 func main() {
 	hostPtr := flag.String("host", "", "Redis Server FQDN:port")
-	//usernamePtr := flag.String("username", "", "Redis Server username")
-	//passwordPtr := flag.String("password", "", "Redis user password")
+	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")
+	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
-	//todo add username/password support
-	//username := *usernamePtr
-	//password := *passwordPtr
+	username := *usernamePtr
+	password := *passwordPtr
 	//db := *dbPtr
 	initializeDB := *initializeDBPtr
 	rate := *ratePtr
@@ -216,11 +220,11 @@ func main() {
 	durationInt := *testDurationPtr
 	duration := time.Second * time.Duration(durationInt)
 
-	pool := newPool(host)
+	pool := newPool(host, username, password)
 	client := pool.Get()
 	defer client.Close()
 
-	if !initializeDB {
+	if initializeDB {
 		buildTestData(&client, 50000)
 	}
 	getRateTest(&client, rate, keyMaxValue, duration)