Redis ID generator benchmark

go.mod

module ridgen

go 1.20

require (
	github.com/go-redis/redis/v8 v8.11.5
	golang.org/x/sync v0.2.0
)

require (
	github.com/cespare/xxhash/v2 v2.1.2 // indirect
	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
)

idgen_test.go

package main

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/go-redis/redis/v8"
	"golang.org/x/sync/errgroup"
)

func BenchmarkWithWatch(b *testing.B) {
	concurrencies := []int{1, 10, 100, 1000}
	for _, c := range concurrencies {
		b.Run(fmt.Sprintf("Concurrency:%d", c), func(b *testing.B) {
			if err := loop(b.N, c, AcquireNewIDWithWatch); err != nil {
				b.Fatal(err)
			}
		})
	}
}

func BenchmarkWithLua(b *testing.B) {
	concurrencies := []int{1, 10, 100, 1000}
	for _, c := range concurrencies {
		b.Run(fmt.Sprintf("Concurrency:%d", c), func(b *testing.B) {
			if err := loop(b.N, c, AcquireNewIDWithLua); err != nil {
				b.Fatal(err)
			}
		})
	}
}

func loop(ids, concurrency int, f func(context.Context, *IDManager, int) (int64, error)) error {
	eg, ctx := errgroup.WithContext(context.Background())
	eg.SetLimit(concurrency)

	client := redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		PoolSize: concurrency,
	})
	m := &IDManager{
		redis:        client,
		IDCounterKey: "item_id",
	}
	if _, err := m.InitNewID(ctx, 1, time.Hour); err != nil {
		return err
	}

	for i := 0; i < ids; i++ {
		loop := i
		eg.Go(func() error {
			_, err := f(ctx, m, loop)
			if err != nil {
				return err
			}
			return nil
		})
	}

	if err := eg.Wait(); err != nil {
		return err
	}
	return nil
}

main.go

package main

import (
	"context"
	"errors"
	"math/rand"
	"strings"
	"time"

	"github.com/go-redis/redis/v8"
)

func AcquireNewIDWithWatch(ctx context.Context, m *IDManager, loop int) (int64, error) {
	var newID int64
	var acquireErr error
	ok := false
	// 1000回までリトライする
	for i := 0; i < 1000; i++ {
		select {
		case <-ctx.Done():
			break
		default:
		}
		// 再試行の場合は1-5msのランダムな待ち時間を入れる
		if i > 0 {
			time.Sleep(time.Millisecond * time.Duration(1+rand.Intn(4)))
		}
		id, exists, err := m.AcquireNewIDWithWatch(ctx, time.Hour)
		if err != nil {
			acquireErr = err
			continue
		}
		if !exists {
			return 0, errors.New("key does not exist")
		}

		// 成功
		newID = id
		ok = true
		acquireErr = nil
		break
	}

	if ok {
		return newID, nil
	}
	if acquireErr != nil {
		return 0, acquireErr
	}
	return 0, errors.New("failed to acquire new id")
}

func AcquireNewIDWithLua(ctx context.Context, m *IDManager, loop int) (int64, error) {
	id, exists, err := m.AcquireNewIDWithLua(ctx, time.Hour)
	if err != nil {
		return 0, err
	}
	if !exists {
		return 0, errors.New("key does not exist")
	}
	return id, nil
}

type IDManager struct {
	redis        *redis.Client
	IDCounterKey string
}

func (m *IDManager) AcquireNewIDWithWatch(ctx context.Context, updateTTL time.Duration) (int64, bool, error) {
	var keyExists bool
	var newID int64
	err := m.redis.Watch(ctx, func(tx *redis.Tx) error {
		txExists := tx.Exists(ctx, m.IDCounterKey)
		exists, err := txExists.Result()
		if err != nil {
			return err
		}
		if exists > 0 {
			keyExists = true
		} else {
			keyExists = false
			return nil
		}

		var incrCmd *redis.IntCmd
		_, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
			incrCmd = pipe.Incr(ctx, m.IDCounterKey)
			_ = pipe.Expire(ctx, m.IDCounterKey, updateTTL)
			return nil
		})
		if err != nil {
			return err
		}
		newID = incrCmd.Val()
		return nil
	}, m.IDCounterKey)

	return newID, keyExists, err
}

func (m *IDManager) AcquireNewIDWithLua(ctx context.Context, updateTTL time.Duration) (int64, bool, error) {
	script := strings.TrimSpace(`
local e=redis.call("EXISTS", KEYS[1])
if e==0 then
  return {0, 0}
end
local n=redis.call("INCR", KEYS[1])
redis.call("EXPIRE", KEYS[1], ARGV[1])
return {1, n}
	`)

	incr := redis.NewScript(script)
	res, err := incr.Eval(ctx, m.redis, []string{m.IDCounterKey}, updateTTL.Seconds()).Int64Slice()
	if err != nil {
		return 0, false, err
	}
	keyExists := res[0] == 1
	newID := res[1]
	return newID, keyExists, err
}

func (m *IDManager) InitNewID(ctx context.Context, newID int64, ttl time.Duration) (bool, error) {
	script := strings.TrimSpace(`
local e=redis.call("EXISTS", KEYS[1])
if e>0 then
  return 0
end
redis.call("SET", KEYS[1], ARGV[1], "EX", ARGV[2])
return 1
	`)

	init := redis.NewScript(script)
	res, err := init.Eval(ctx, m.redis, []string{m.IDCounterKey}, newID, ttl.Seconds()).Int()
	if err != nil {
		return false, err
	}
	ok := res == 1
	return ok, nil
}