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 } }