package cache import ( "context" "errors" "time" "github.com/EthanCodeCraft/xlgo-core/database" "github.com/EthanCodeCraft/xlgo-core/utils" ) // 分布式锁错误 var ( ErrLockNotHeld = errors.New("锁未被当前客户端持有") ErrLockExpired = errors.New("锁已过期") ErrRedisNotReady = errors.New("Redis 未初始化") ) // LockToken 锁令牌(用于安全释放锁) 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) { if database.RedisClient == nil { return nil, ErrRedisNotReady } token := utils.UUID() ttlMs := int64(ttl / time.Millisecond) result, err := database.RedisClient.Eval(ctx, lockScript, []string{key}, token, ttlMs).Result() if err != nil { return nil, err } if result.(int64) == 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 { if database.RedisClient == nil { return ErrRedisNotReady } if token == nil { return ErrLockNotHeld } result, err := database.RedisClient.Eval(ctx, unlockScript, []string{token.Key}, token.Token).Result() if err != nil { return err } if result.(int64) == 0 { return ErrLockNotHeld } return nil } // UnlockByKey 按键名释放锁(不安全,仅用于旧代码兼容) // 注意: 此函数不检查 Token,任何客户端都能释放锁 func UnlockByKey(ctx context.Context, key string) error { if database.RedisClient == nil { return nil } return database.RedisClient.Del(ctx, key).Err() } // ExtendLock 续期锁 // 参数: token 锁令牌,ttl 新的过期时间 func ExtendLock(ctx context.Context, token *LockToken, ttl time.Duration) error { if database.RedisClient == nil { return ErrRedisNotReady } if token == nil { return ErrLockNotHeld } ttlMs := int64(ttl / time.Millisecond) result, err := database.RedisClient.Eval(ctx, extendScript, []string{token.Key}, token.Token, ttlMs).Result() if err != nil { return err } if result.(int64) == 0 { return ErrLockNotHeld } return nil } // TryLock 尝试获取锁,失败时等待重试 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 } time.Sleep(retryInterval) } return nil, nil } // WithLock 使用分布式锁执行函数(自动管理锁) // 参数: key 锁名称,ttl 锁定时长,fn 业务函数 // 注意: 如果任务执行时间超过 ttl,需要设置更长的 ttl 或使用 WithLockAutoExtend 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 Unlock(ctx, token) return fn() } // WithLockAutoExtend 使用分布式锁执行函数(自动续期) // 参数: key 锁名称,initialTTL 初始锁定时长,extendInterval 续期间隔,fn 业务函数 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 // 未获取到锁,跳过执行 } // 启动续期协程 done := make(chan struct{}) go func() { defer close(done) ticker := time.NewTicker(extendInterval) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-done: return case <-ticker.C: // 续期锁(每次续期为 initialTTL) if err := ExtendLock(ctx, token, initialTTL); err != nil { return // 续期失败,停止续期 } } } }() // 执行业务函数 err = fn() // 停止续期并释放锁 done <- struct{}{} Unlock(ctx, token) return err } // IsLocked 检查锁是否被占用(不获取锁) func IsLocked(ctx context.Context, key string) (bool, error) { if database.RedisClient == nil { return false, nil } return database.RedisClient.Exists(ctx, key).Val() > 0, nil } // GetLockTTL 获取锁的剩余过期时间 func GetLockTTL(ctx context.Context, key string) (time.Duration, error) { if database.RedisClient == nil { return 0, nil } return database.RedisClient.TTL(ctx, key).Result() } // ForceUnlock 强制释放锁(危险操作,仅用于管理场景) // 注意: 此函数不检查 Token,强制删除锁 func ForceUnlock(ctx context.Context, key string) error { if database.RedisClient == nil { return nil } return database.RedisClient.Del(ctx, key).Err() } // ===== 计数器操作 ===== // Incr 自增计数器 func Incr(ctx context.Context, key string) (int64, error) { if database.RedisClient == nil { return 0, nil } return database.RedisClient.Incr(ctx, key).Result() } // IncrBy 指定增量自增 func IncrBy(ctx context.Context, key string, value int64) (int64, error) { if database.RedisClient == nil { return 0, nil } return database.RedisClient.IncrBy(ctx, key, value).Result() } // Decr 自减计数器 func Decr(ctx context.Context, key string) (int64, error) { if database.RedisClient == nil { return 0, nil } return database.RedisClient.Decr(ctx, key).Result() } // GetTTL 获取键的剩余过期时间 func GetTTL(ctx context.Context, key string) (time.Duration, error) { if database.RedisClient == nil { return 0, nil } return database.RedisClient.TTL(ctx, key).Result() } // SetExpire 设置键的过期时间 func SetExpire(ctx context.Context, key string, ttl time.Duration) (bool, error) { if database.RedisClient == nil { return false, nil } return database.RedisClient.Expire(ctx, key, ttl).Result() } // GetRaw 获取原始字符串值(不反序列化) func GetRaw(ctx context.Context, key string) (string, error) { if database.RedisClient == nil { return "", nil } return database.RedisClient.Get(ctx, key).Result() } // SetRaw 设置原始值(不序列化) func SetRaw(ctx context.Context, key string, value string, ttl time.Duration) error { if database.RedisClient == nil { return nil } return database.RedisClient.Set(ctx, key, value, ttl).Err() }