114 lines
2.7 KiB
Go
114 lines
2.7 KiB
Go
|
package lock
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/redis/go-redis/v9"
|
||
|
|
)
|
||
|
|
|
||
|
|
var (
|
||
|
|
ErrLockNotAcquired = errors.New("lock not acquired")
|
||
|
|
ErrLockNotHeld = errors.New("lock not held by this owner")
|
||
|
|
)
|
||
|
|
|
||
|
|
// DistributedLock 基于 Redis 的分布式锁
|
||
|
|
type DistributedLock struct {
|
||
|
|
client *redis.Client
|
||
|
|
key string
|
||
|
|
owner string // 唯一标识，用于安全释放锁
|
||
|
|
ttl time.Duration
|
||
|
|
}
|
||
|
|
|
||
|
|
// New 创建一个分布式锁实例
|
||
|
|
// - client: Redis 客户端
|
||
|
|
// - key: 锁的 key
|
||
|
|
// - owner: 持有者唯一标识（如 UUID）
|
||
|
|
// - ttl: 锁的过期时间，防止死锁
|
||
|
|
func New(client *redis.Client, key string, owner string, ttl time.Duration) *DistributedLock {
|
||
|
|
return &DistributedLock{
|
||
|
|
client: client,
|
||
|
|
key: key,
|
||
|
|
owner: owner,
|
||
|
|
ttl: ttl,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Acquire 尝试获取锁（非阻塞）
|
||
|
|
func (l *DistributedLock) Acquire(ctx context.Context) (bool, error) {
|
||
|
|
ok, err := l.client.SetNX(ctx, l.key, l.owner, l.ttl).Result()
|
||
|
|
if err != nil {
|
||
|
|
return false, err
|
||
|
|
}
|
||
|
|
return ok, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// AcquireWithRetry 带重试的获取锁
|
||
|
|
// - retryInterval: 重试间隔
|
||
|
|
// - maxRetries: 最大重试次数，0 表示无限重试
|
||
|
|
func (l *DistributedLock) AcquireWithRetry(ctx context.Context, retryInterval time.Duration, maxRetries int) (bool, error) {
|
||
|
|
for i := 0; maxRetries == 0 || i < maxRetries; i++ {
|
||
|
|
acquired, err := l.Acquire(ctx)
|
||
|
|
if err != nil {
|
||
|
|
return false, err
|
||
|
|
}
|
||
|
|
if acquired {
|
||
|
|
return true, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
select {
|
||
|
|
case <-ctx.Done():
|
||
|
|
return false, ctx.Err()
|
||
|
|
case <-time.After(retryInterval):
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return false, ErrLockNotAcquired
|
||
|
|
}
|
||
|
|
|
||
|
|
// Release 释放锁（仅当自己持有时才释放，使用 Lua 脚本保证原子性）
|
||
|
|
func (l *DistributedLock) Release(ctx context.Context) error {
|
||
|
|
script := redis.NewScript(`
|
||
|
|
if redis.call("GET", KEYS[1]) == ARGV[1] then
|
||
|
|
return redis.call("DEL", KEYS[1])
|
||
|
|
else
|
||
|
|
return 0
|
||
|
|
end
|
||
|
|
`)
|
||
|
|
|
||
|
|
result, err := script.Run(ctx, l.client, []string{l.key}, l.owner).Int64()
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if result == 0 {
|
||
|
|
return ErrLockNotHeld
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Refresh 续期锁（延长 TTL，仅当自己持有时才续期）
|
||
|
|
func (l *DistributedLock) Refresh(ctx context.Context) (bool, error) {
|
||
|
|
script := redis.NewScript(`
|
||
|
|
if redis.call("GET", KEYS[1]) == ARGV[1] then
|
||
|
|
return redis.call("PEXPIRE", KEYS[1], ARGV[2])
|
||
|
|
else
|
||
|
|
return 0
|
||
|
|
end
|
||
|
|
`)
|
||
|
|
|
||
|
|
result, err := script.Run(ctx, l.client, []string{l.key}, l.owner, l.ttl.Milliseconds()).Int64()
|
||
|
|
if err != nil {
|
||
|
|
return false, err
|
||
|
|
}
|
||
|
|
return result == 1, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// IsHeld 检查锁是否被持有（不一定是自己）
|
||
|
|
func (l *DistributedLock) IsHeld(ctx context.Context) (bool, error) {
|
||
|
|
val, err := l.client.Exists(ctx, l.key).Result()
|
||
|
|
if err != nil {
|
||
|
|
return false, err
|
||
|
|
}
|
||
|
|
return val > 0, nil
|
||
|
|
}
|