Skip to content

Instantly share code, notes, and snippets.

@maguec
Created March 4, 2022 21:31
Show Gist options
  • Save maguec/f308d0c3d4f09bbdd4c3d48b0b63a1ba to your computer and use it in GitHub Desktop.
Save maguec/f308d0c3d4f09bbdd4c3d48b0b63a1ba to your computer and use it in GitHub Desktop.
Redis pub/sub benchmark publisher
package main
import (
"context"
"crypto/sha256"
"fmt"
"log"
"os"
"strconv"
"time"
"github.com/go-redis/redis/v8"
"github.com/pborman/getopt/v2"
"go.uber.org/ratelimit"
)
func worker(
id int,
ctx context.Context,
jobs <-chan int,
results chan<- int,
redisClient *redis.Client,
rl ratelimit.Limiter) {
h := sha256.New()
for j := range jobs {
rl.Take()
num := strconv.Itoa(j)
h.Write([]byte(num))
key := fmt.Sprintf("%x", h.Sum(nil))
err := redisClient.Publish(ctx, "KEY", key).Err()
if err != nil {
log.Panic(err)
}
results <- j
}
}
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', 10000, "limit to this number of RPS")
getopt.Parse()
if *helpFlag {
getopt.PrintUsage(os.Stderr)
os.Exit(1)
}
client := redis.NewClient(&redis.Options{
Addr: fmt.Sprintf("%s:%d", *redisHost, *redisPort),
Password: *redisPassword,
PoolSize: *threadCount,
MinIdleConns: *threadCount,
PoolTimeout: 0,
IdleTimeout: 20 * time.Second,
DialTimeout: 2 * time.Second,
})
jobs := make(chan int, *messageCount)
results := make(chan int, *messageCount)
rl := ratelimit.New(*rps)
for w := 1; w <= *threadCount; w++ {
go worker(w, ctx, jobs, results, client, rl)
}
for j := 0; j <= *messageCount-1; j++ {
jobs <- j
}
close(jobs)
// Finally we collect all the results of the work.
for a := 0; a <= *messageCount-1; a++ {
<-results
}
os.Exit(0)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment