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

package redisx

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

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

// 分布式锁

// errLockNotOwner 表示当前实例已不再持有该锁（锁已过期或被他人抢占）。
// 用哨兵错误以便 errors.Is 判定，供观测日志区分“锁被抢走”与其他失败。
var errLockNotOwner = errors.New("lock not owner")

type RedisLock struct {
	key      string
	token    string
	ttl      time.Duration
	acquired time.Time // 加锁成功的时刻，用于统计持锁时长（配合 time.Since 使用，单调）
	label    string    // 业务标签，从 key 推断，便于按锁类型聚合观测日志
}

// lockLabel 从 key 推断业务标签，便于日志聚合。
// 约定 key 形如 "<domain>:<category>:<id>"，取前两段，例如
// "data-download:content:xxx" -> "data-download:content"。无法识别时返回 "unknown"。
func lockLabel(key string) string {
	parts := strings.SplitN(key, ":", 3)
	if len(parts) >= 2 {
		return parts[0] + ":" + parts[1]
	}
	if len(parts) == 1 && parts[0] != "" {
		return parts[0]
	}
	return "unknown"
}

// 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, acquired: time.Now(), label: lockLabel(key)}, true, nil
}

// TryLockWithToken 使用调用方指定的 token 非阻塞加锁。
// 若锁已被同一 token 持有（例如进程重启后同一任务接管自己留下的锁），视为加锁成功并续期。
func TryLockWithToken(ctx context.Context, key, token string, ttl time.Duration) (*RedisLock, bool, error) {
	const luaScript = `
if redis.call("set", KEYS[1], ARGV[1], "NX", "PX", ARGV[2]) then
	return 1
end
if redis.call("get", KEYS[1]) == ARGV[1] then
	redis.call("pexpire", KEYS[1], ARGV[2])
	return 1
end
return 0
`
	res, err := redis.NewScript(luaScript).Run(ctx, getClient(), []string{key}, token, ttl.Milliseconds()).Int()
	if err != nil {
		return nil, false, err
	}
	if res == 0 {
		return nil, false, nil
	}
	return &RedisLock{key: key, token: token, ttl: ttl, acquired: time.Now(), label: lockLabel(key)}, 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 fmt.Errorf("unlock failed: %w", errLockNotOwner)
	}
	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 fmt.Errorf("unlock failed: %w", errLockNotOwner)
	}
	return nil
}

// SafeRelease 带日志保护版本。
//
// 除了尽力释放锁外，它还承担 DD-001 的观测职责：
//   - 若释放时发现锁已不属于自己（"not owner"），说明临界区耗时超过 TTL、
//     锁在持有期间被别的持有者抢走，可能已发生并发覆盖，记为 Error（关键信号）。
//   - 若正常释放但持锁时长已逼近 TTL（>70%），记为 Warn（风险预警）。
//   - 其他释放错误记为 Error。
// 该方法不改变控制流，也不返回错误，调用方无感知。
func (l *RedisLock) SafeRelease(ctx context.Context) {
	held := time.Since(l.acquired)
	err := l.Release(ctx)
	switch {
	case err == nil:
		if l.ttl > 0 && held > l.ttl*7/10 {
			zap.L().Warn("redis lock held close to ttl",
				zap.String("label", l.label),
				zap.String("key", l.key),
				zap.Duration("held", held),
				zap.Duration("ttl", l.ttl))
		}
	case errors.Is(err, errLockNotOwner):
		// 关键信号：干完活时锁已不是自己的，临界区超时被抢，可能已并发覆盖。
		zap.L().Error("redis lock lost before release (possible concurrent overwrite)",
			zap.String("label", l.label),
			zap.String("key", l.key),
			zap.String("token", l.token),
			zap.Duration("held", held),
			zap.Duration("ttl", l.ttl))
	default:
		zap.L().Error("redis lock release failed",
			zap.String("label", l.label),
			zap.String("key", l.key),
			zap.Duration("held", held),
			zap.Error(err))
	}
}

// Extend renews the lock only while this instance still owns it.
func (l *RedisLock) Extend(ctx context.Context, ttl time.Duration) error {
	const luaScript = `
if redis.call("get", KEYS[1]) == ARGV[1] then
	return redis.call("pexpire", KEYS[1], ARGV[2])
else
	return 0
end
`
	res, err := redis.NewScript(luaScript).Run(ctx, getClient(), []string{l.key}, l.token, ttl.Milliseconds()).Int()
	if err != nil {
		return err
	}
	if res == 0 {
		return fmt.Errorf("extend failed: %w", errLockNotOwner)
	}
	l.ttl = ttl
	return nil
}

// IsOwner reports whether this instance still owns the lock.
func (l *RedisLock) IsOwner(ctx context.Context) (bool, error) {
	value, err := getClient().Get(ctx, l.key).Result()
	if err == redis.Nil {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return value == l.token, nil
}

// SetIfOwner writes a completion marker only while this instance still owns
// the lock. It is used by lease-based workflows to avoid publishing a stale
// worker's result after its lease has expired.
func (l *RedisLock) SetIfOwner(ctx context.Context, key, value string, ttl time.Duration) error {
	const luaScript = `
if redis.call("get", KEYS[1]) == ARGV[1] then
	redis.call("set", KEYS[2], ARGV[2], "PX", ARGV[3])
	return 1
else
	return 0
end
`
	res, err := redis.NewScript(luaScript).Run(ctx, getClient(), []string{l.key, key}, l.token, value, ttl.Milliseconds()).Int()
	if err != nil {
		return err
	}
	if res == 0 {
		return fmt.Errorf("commit failed: %w", errLockNotOwner)
	}
	return nil
}

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

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