RateLimitedSetex

gistfile1.txt

package main

import (
	"context"
	"crypto/sha256"
	"fmt"
	"log"
	"os"
	"strconv"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/jamiealquiza/tachymeter"
	"github.com/pborman/getopt/v2"
	"go.uber.org/ratelimit"
)

type redisCall struct {
	key     string
	num     int
	unixts  int64
	elapsed time.Duration
}

func worker(id int, ctx context.Context, jobs <-chan int, results chan<- redisCall, redisClient *redis.ClusterClient, rps int) {
	h := sha256.New()
	rl := ratelimit.New(rps)
	for j := range jobs {
		rl.Take()
		num := strconv.Itoa(j)
		h.Write([]byte(num))
		key := fmt.Sprintf("%x", h.Sum(nil))
		startTime := time.Now()
		err := redisClient.SetEX(ctx, key, "XXXXXXXXXXXXXXXXXXXXXXXXXX", 30*time.Second).Err()
		if err != nil {
			log.Panic(err)
		}
		results <- redisCall{key: key, num: j, unixts: startTime.Unix(), elapsed: time.Since(startTime)}
	}
}

func main() {

	var ctx = context.Background()
	helpFlag := getopt.BoolLong("help", 'h', "display help")

	redisHost := getopt.StringLong("host", 's', "127.0.0.1", "Redis Host")
	redisPassword := getopt.StringLong("password", 'a', "", "Redis Password")
	redisPort := getopt.IntLong("port", 'p', 6379, "Redis Port")
	messageCount := getopt.IntLong("message-count", 'c', 100000, "run this man times")
	threadCount := getopt.IntLong("threads", 't', 1, "run this many threads")
	rps := getopt.IntLong("requests-per-second", 'r', 10, "limit each thread to this number of RPS")

	getopt.Parse()

	if *helpFlag {
		getopt.PrintUsage(os.Stderr)
		os.Exit(1)
	}

	client := redis.NewClusterClient(&redis.ClusterOptions{
		Addrs:        []string{fmt.Sprintf("%s:%d", *redisHost, *redisPort)},
		MaxRedirects: 3,
		Password:     *redisPassword,
		PoolSize:     *threadCount,
		MinIdleConns: *threadCount,
		PoolTimeout:  0,
		IdleTimeout:  20 * time.Second,
		DialTimeout:  2 * time.Second,
	})

	jobs := make(chan int, *messageCount)
	results := make(chan redisCall, *messageCount)

	for w := 1; w <= *threadCount; w++ {
		go worker(w, ctx, jobs, results, client, *rps)
	}

	for j := 0; j <= *messageCount-1; j++ {
		jobs <- j
	}
	close(jobs)

	t := tachymeter.New(&tachymeter.Config{Size: *messageCount})
	// Finally we collect all the results of the work.
	for a := 0; a <= *messageCount-1; a++ {
		r := <-results
		//fmt.Printf("%+v\n", r)
		t.AddTime(r.elapsed)
	}
	fmt.Println(t.Calc())
	os.Exit(0)

}