FranklinWH-Toolbox/service/cache/redisx/lockx.go

package redisx

import (
	"context"
	"errors"
	"time"

	"github.com/google/uuid"
	"github.com/redis/go-redis/v9"
)

// 分布式锁

type RedisLock struct {
	key   string
	token string
	ttl   time.Duration
}

// TryLock 非阻塞尝试加锁
func TryLock(ctx context.Context, key string, ttl time.Duration) (*RedisLock, bool, error) {
	token := uuid.NewString()
	ok, err := getClient().SetNX(ctx, key, token, ttl).Result()
	if err != nil {
		return nil, false, err
	}
	if !ok {
		return nil, false, nil
	}
	return &RedisLock{key: key, token: token, ttl: ttl}, true, nil
}

// Lock 阻塞加锁，可设置重试间隔和超时
func Lock(ctx context.Context, key string, ttl, retryInterval, timeout time.Duration) (*RedisLock, error) {
	start := time.Now()
	for {
		lock, ok, err := TryLock(ctx, key, ttl)
		if err != nil {
			return nil, err
		}
		if ok {
			return lock, nil
		}
		if timeout > 0 && time.Since(start) >= timeout {
			return nil, errors.New("lock timeout")
		}
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(retryInterval):
		}
	}
}

// Release 安全释放锁（Lua 脚本） - 通过key和token直接释放
func LockRelease(ctx context.Context, key string, token string) error {
	const luaScript = `
if redis.call("get", KEYS[1]) == ARGV[1] then
	return redis.call("del", KEYS[1])
else
	return 0
end
`
	res, err := redis.NewScript(luaScript).Run(ctx, getClient(), []string{key}, token).Int()
	if err != nil {
		return err
	}
	if res == 0 {
		return errors.New("unlock failed: not owner")
	}
	return nil
}

// Release 安全释放锁（Lua 脚本）
func (l *RedisLock) Release(ctx context.Context) error {
	const luaScript = `
if redis.call("get", KEYS[1]) == ARGV[1] then
	return redis.call("del", KEYS[1])
else
	return 0
end
`
	res, err := redis.NewScript(luaScript).Run(ctx, getClient(), []string{l.key}, l.token).Int()
	if err != nil {
		return err
	}
	if res == 0 {
		return errors.New("unlock failed: not owner")
	}
	return nil
}

// SafeRelease 带日志保护版本
func (l *RedisLock) SafeRelease(ctx context.Context) {
	_ = l.Release(ctx)
}

func (l *RedisLock) GetKey() string {
	return l.key
}

func (l *RedisLock) GetToken() string {
	return l.token
}