thread-master/apps/backend/internal/lib/redislock/redislock.go

package redislock

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/zeromicro/go-zero/core/stores/redis"
)

const (
	renewScript = `
if redis.call('GET', KEYS[1]) == ARGV[1] then
  return redis.call('PEXPIRE', KEYS[1], ARGV[2])
end
return 0
`
	releaseScript = `
if redis.call('GET', KEYS[1]) == ARGV[1] then
  return redis.call('DEL', KEYS[1])
end
return 0
`
)

// Client is the subset of Redis used by Lock.
type Client interface {
	SetnxExCtx(ctx context.Context, key, value string, seconds int) (bool, error)
	EvalCtx(ctx context.Context, script string, keys []string, args ...any) (any, error)
}

// Lock is a small Redis lease. Each acquisition gets a unique token so an old
// owner cannot renew or release a successor's lease.
type Lock struct {
	client Client
	key    string
	prefix string
	ttl    time.Duration

	mu    sync.RWMutex
	token string
}

func New(conf redis.RedisConf, key, ownerPrefix string, ttl time.Duration) *Lock {
	return NewWithClient(redis.MustNewRedis(conf), key, ownerPrefix, ttl)
}

func NewWithClient(client Client, key, ownerPrefix string, ttl time.Duration) *Lock {
	return &Lock{client: client, key: key, prefix: strings.TrimSpace(ownerPrefix), ttl: ttl}
}

func (l *Lock) Acquire(ctx context.Context) (bool, error) {
	if err := l.validate(); err != nil {
		return false, err
	}
	token, err := newToken(l.prefix)
	if err != nil {
		return false, err
	}
	acquired, err := l.client.SetnxExCtx(ctx, l.key, token, durationSeconds(l.ttl))
	if err != nil || !acquired {
		return acquired, err
	}
	l.mu.Lock()
	l.token = token
	l.mu.Unlock()
	return true, nil
}

// Renew extends the lease only while this acquisition still owns it.
func (l *Lock) Renew(ctx context.Context) (bool, error) {
	if err := l.validate(); err != nil {
		return false, err
	}
	token := l.currentToken()
	if token == "" {
		return false, nil
	}
	ttlMilliseconds := l.ttl.Milliseconds()
	if ttlMilliseconds < 1 {
		ttlMilliseconds = 1
	}
	result, err := l.client.EvalCtx(ctx, renewScript, []string{l.key}, token, ttlMilliseconds)
	if err != nil {
		return false, err
	}
	return resultIsOne(result), nil
}

func (l *Lock) Release(ctx context.Context) error {
	if err := l.validate(); err != nil {
		return err
	}
	token := l.currentToken()
	if token == "" {
		return nil
	}
	result, err := l.client.EvalCtx(ctx, releaseScript, []string{l.key}, token)
	if err == nil && resultIsOne(result) {
		l.mu.Lock()
		if l.token == token {
			l.token = ""
		}
		l.mu.Unlock()
	}
	return err
}

func (l *Lock) validate() error {
	if l == nil || l.client == nil || l.key == "" || l.prefix == "" || l.ttl <= 0 {
		return fmt.Errorf("invalid redis lock")
	}
	return nil
}

func (l *Lock) currentToken() string {
	l.mu.RLock()
	defer l.mu.RUnlock()
	return l.token
}

func newToken(prefix string) (string, error) {
	random := make([]byte, 16)
	if _, err := rand.Read(random); err != nil {
		return "", fmt.Errorf("generate redis lock token: %w", err)
	}
	return prefix + ":" + hex.EncodeToString(random), nil
}

func durationSeconds(ttl time.Duration) int {
	seconds := int(ttl / time.Second)
	if ttl%time.Second != 0 {
		seconds++
	}
	if seconds < 1 {
		seconds = 1
	}
	return seconds
}

func resultIsOne(result any) bool {
	switch value := result.(type) {
	case int64:
		return value == 1
	case int:
		return value == 1
	case uint64:
		return value == 1
	default:
		return false
	}
}