package cache import ( "context" "errors" "fmt" "time" "github.com/EthanCodeCraft/xlgo-core/database" "github.com/EthanCodeCraft/xlgo-core/utils" ) // 分布式锁错误 var ( ErrLockNotHeld = errors.New("锁未被当前客户端持有") ErrLockExpired = errors.New("锁已过期") ErrRedisNotReady = errors.New("Redis 未初始化") // ErrLockUnexpectedResult Lua 脚本返回了非预期的结果类型(C1b:裸类型断言防护)。 ErrLockUnexpectedResult = errors.New("锁脚本返回非预期结果") ) // toInt64 将 Lua 脚本返回值安全断言为 int64(C1b:禁止裸断言 panic)。 // go-redis 对整数返回 int64,但 nil/错误响应下可能为其他类型。 func toInt64(v any) (int64, error) { n, ok := v.(int64) if !ok { return 0, fmt.Errorf("脚本返回类型 %T: %w", v, ErrLockUnexpectedResult) } return n, nil } // LockToken 锁令牌(用于安全释放锁)。 // // 安全说明(C1d 设计局限):Token 是随机 UUID(非单调递增的 fencing token)。 // 本实现基于 Redis SET PX + Lua CAS,保证"持有者才能解锁/续期",但**无法防 TTL 到期后的 // 双 worker 并发**:若 worker A 因 GC/网络停滞超过 TTL,锁过期后 worker B 获得锁, // A 恢复后仍可能写过期数据。完整的 fencing token 防护需:① 用 Redis INCR 生成单调 token, // ② 下游存储层(DB/外部服务)记录已见最大 token 并拒绝旧 token 写入。 // 框架无法单方面保证②,需下游配合,故本类型仅提供 UUID token。对 TTL 到期敏感的场景, // 请确保 ttl >> 业务最长执行时间,或下游实现 fencing token 校验。 type LockToken struct { Key string // 锁的键名 Token string // 锁的唯一标识(UUID) } // lockScript 加锁 Lua 脚本 // 返回: 1 表示成功加锁,0 表示锁已被占用 const lockScript = ` if redis.call("exists", KEYS[1]) == 0 then redis.call("set", KEYS[1], ARGV[1], "PX", ARGV[2]) return 1 else return 0 end ` // unlockScript 解锁 Lua 脯本 // 只有持有正确 Token 的客户端才能解锁 // 返回: 1 表示成功解锁,0 表示 Token 不匹配(锁不属于该客户端) const unlockScript = ` if redis.call("get", KEYS[1]) == ARGV[1] then redis.call("del", KEYS[1]) return 1 else return 0 end ` // extendScript 续期 Lua 脚本 // 只有持有正确 Token 的客户端才能续期 // 返回: 1 表示成功续期,0 表示 Token 不匹配或锁不存在 const extendScript = ` if redis.call("get", KEYS[1]) == ARGV[1] then redis.call("pexpire", KEYS[1], ARGV[2]) return 1 else return 0 end ` // NewLock 创建分布式锁 // 参数: key 锁名称,ttl 锁定时长 // 返回: LockToken 用于后续解锁或续期 func NewLock(ctx context.Context, key string, ttl time.Duration) (*LockToken, error) { rdb := database.GetRedis() if rdb == nil { return nil, ErrRedisNotReady } token := utils.UUID() ttlMs := int64(ttl / time.Millisecond) result, err := rdb.Eval(ctx, lockScript, []string{key}, token, ttlMs).Result() if err != nil { return nil, err } n, err := toInt64(result) if err != nil { return nil, err } if n == 1 { return &LockToken{Key: key, Token: token}, nil } return nil, nil // 锁已被其他客户端持有 } // Lock 简化的加锁函数(返回 bool) // 注意: 使用此函数无法安全释放锁,建议使用 NewLock func Lock(ctx context.Context, key string, ttl time.Duration) (bool, error) { token, err := NewLock(ctx, key, ttl) if err != nil { return false, err } return token != nil, nil } // Unlock 安全释放锁 func Unlock(ctx context.Context, token *LockToken) error { rdb := database.GetRedis() if rdb == nil { return ErrRedisNotReady } if token == nil { return ErrLockNotHeld } result, err := rdb.Eval(ctx, unlockScript, []string{token.Key}, token.Token).Result() if err != nil { return err } n, err := toInt64(result) if err != nil { return err } if n == 0 { return ErrLockNotHeld } return nil } // UnlockByKey 按键名释放锁(不安全,仅用于旧代码兼容) // 注意: 此函数不检查 Token,任何客户端都能释放锁 func UnlockByKey(ctx context.Context, key string) error { rdb := database.GetRedis() if rdb == nil { return nil } return rdb.Del(ctx, key).Err() } // ExtendLock 续期锁 // 参数: token 锁令牌,ttl 新的过期时间 func ExtendLock(ctx context.Context, token *LockToken, ttl time.Duration) error { rdb := database.GetRedis() if rdb == nil { return ErrRedisNotReady } if token == nil { return ErrLockNotHeld } ttlMs := int64(ttl / time.Millisecond) result, err := rdb.Eval(ctx, extendScript, []string{token.Key}, token.Token, ttlMs).Result() if err != nil { return err } n, err := toInt64(result) if err != nil { return err } if n == 0 { return ErrLockNotHeld } return nil } // TryLock 尝试获取锁,失败时等待重试。重试等待响应 ctx 取消(C1c 修复)。 func TryLock(ctx context.Context, key string, ttl time.Duration, retryInterval time.Duration, maxRetry int) (*LockToken, error) { for i := 0; i < maxRetry; i++ { token, err := NewLock(ctx, key, ttl) if err != nil { return nil, err } if token != nil { return token, nil } // 响应 ctx 取消,避免最长阻塞 maxRetry*retryInterval(C1c:禁止 time.Sleep 无视 ctx)。 select { case <-ctx.Done(): return nil, ctx.Err() case <-time.After(retryInterval): } } return nil, nil } // WithLock 使用分布式锁执行函数(自动管理锁)。 // 参数: key 锁名称,ttl 锁定时长,fn 业务函数 // 注意: 如果任务执行时间超过 ttl,需要设置更长的 ttl 或使用 WithLockAutoExtend。 // // 解锁用独立 Background ctx(C1a 一致性修复):fn 返回或 panic 后,原 ctx 可能已被 // 调用方取消,用其解锁会失败导致锁泄漏到 TTL。fn panic 时 defer 也保证解锁执行。 func WithLock(ctx context.Context, key string, ttl time.Duration, fn func() error) error { token, err := NewLock(ctx, key, ttl) if err != nil { return err } if token == nil { return nil // 未获取到锁,跳过执行 } defer func() { unlockCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _ = Unlock(unlockCtx, token) }() return fn() } // WithLockAutoExtend 使用分布式锁执行函数(自动续期)。 // 参数: key 锁名称,initialTTL 初始锁定时长,extendInterval 续期间隔,fn 业务函数 // // 并发安全说明(C1a 修复):续期 goroutine 与父用"父关停 + 子 ack"双 channel 协调—— // 父用 close(stop) 通知子退出(close 由唯一所有者执行,安全),子用 close(finished) ack。 // 避免旧实现 done 无缓冲 + 子 defer close(done) + 父 done<-struct{}{} 的 send-on-closed panic // (ctx 取消或 ExtendLock 失败时 done 已 closed,父再 send 即 panic,Unlock 不执行、锁泄漏到 TTL)。 // Unlock 用 context.Background() 派生超时,避免原 ctx 已取消致 Unlock 失败再泄漏。 func WithLockAutoExtend(ctx context.Context, key string, initialTTL time.Duration, extendInterval time.Duration, fn func() error) error { token, err := NewLock(ctx, key, initialTTL) if err != nil { return err } if token == nil { return nil // 未获取到锁,跳过执行 } // 父关停信号(仅父 close)与子 ack 信号(仅子 close)。 stop := make(chan struct{}) finished := make(chan struct{}) go func() { defer close(finished) // 子退出时 ack,父等待 finished ticker := time.NewTicker(extendInterval) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-stop: return case <-ticker.C: // 续期锁(每次续期为 initialTTL)。续期失败则停止续期,fn 应尽快结束。 if err := ExtendLock(ctx, token, initialTTL); err != nil { return } } } }() // defer 兜底:fn panic 时也要停止续期 goroutine 并释放锁(C1a panic 路径修复)。 // 无 defer 时 fn panic 会导致 close(stop) 不执行 → 续期 goroutine 永久泄漏,且 Unlock 不执行 → 锁泄漏到 TTL。 defer func() { close(stop) <-finished unlockCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _ = Unlock(unlockCtx, token) }() // 执行业务函数。 err = fn() return err } // IsLocked 检查锁是否被占用(不获取锁) func IsLocked(ctx context.Context, key string) (bool, error) { rdb := database.GetRedis() if rdb == nil { return false, nil } return rdb.Exists(ctx, key).Val() > 0, nil } // GetLockTTL 获取锁的剩余过期时间 func GetLockTTL(ctx context.Context, key string) (time.Duration, error) { rdb := database.GetRedis() if rdb == nil { return 0, nil } return rdb.TTL(ctx, key).Result() } // ForceUnlock 强制释放锁(危险操作,仅用于管理场景) // 注意: 此函数不检查 Token,强制删除锁 func ForceUnlock(ctx context.Context, key string) error { rdb := database.GetRedis() if rdb == nil { return nil } return rdb.Del(ctx, key).Err() } // ===== 计数器操作 ===== // Incr 自增计数器 func Incr(ctx context.Context, key string) (int64, error) { rdb := database.GetRedis() if rdb == nil { return 0, nil } return rdb.Incr(ctx, key).Result() } // IncrBy 指定增量自增 func IncrBy(ctx context.Context, key string, value int64) (int64, error) { rdb := database.GetRedis() if rdb == nil { return 0, nil } return rdb.IncrBy(ctx, key, value).Result() } // Decr 自减计数器 func Decr(ctx context.Context, key string) (int64, error) { rdb := database.GetRedis() if rdb == nil { return 0, nil } return rdb.Decr(ctx, key).Result() } // GetTTL 获取键的剩余过期时间 func GetTTL(ctx context.Context, key string) (time.Duration, error) { rdb := database.GetRedis() if rdb == nil { return 0, nil } return rdb.TTL(ctx, key).Result() } // SetExpire 设置键的过期时间 func SetExpire(ctx context.Context, key string, ttl time.Duration) (bool, error) { rdb := database.GetRedis() if rdb == nil { return false, nil } return rdb.Expire(ctx, key, ttl).Result() } // GetRaw 获取原始字符串值(不反序列化) func GetRaw(ctx context.Context, key string) (string, error) { rdb := database.GetRedis() if rdb == nil { return "", nil } return rdb.Get(ctx, key).Result() } // SetRaw 设置原始值(不序列化) func SetRaw(ctx context.Context, key string, value string, ttl time.Duration) error { rdb := database.GetRedis() if rdb == nil { return nil } return rdb.Set(ctx, key, value, ttl).Err() }