51 lines
1.3 KiB
Go
51 lines
1.3 KiB
Go
// Package stocklock 演示 Redis 的三个高频用法：分布式锁、滑动窗口限流、库存预扣。
|
|
package stocklock
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"errors"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
var ErrLockNotAcquired = errors.New("没有拿到锁")
|
|
|
|
// 只有持有者（令牌一致）才能释放，防止误删别人的锁
|
|
var unlockScript = redis.NewScript(`
|
|
if redis.call("GET", KEYS[1]) == ARGV[1] then
|
|
return redis.call("DEL", KEYS[1])
|
|
end
|
|
return 0`)
|
|
|
|
type Lock struct {
|
|
rdb *redis.Client
|
|
key string
|
|
token string
|
|
}
|
|
|
|
// TryLock 用 SET key token NX PX ttl 加锁。ttl 必须大于业务最长执行时间（或者自己续期）。
|
|
func TryLock(ctx context.Context, rdb *redis.Client, key string, ttl time.Duration) (*Lock, error) {
|
|
b := make([]byte, 16)
|
|
if _, err := rand.Read(b); err != nil {
|
|
return nil, err
|
|
}
|
|
token := hex.EncodeToString(b)
|
|
ok, err := rdb.SetNX(ctx, key, token, ttl).Result()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if !ok {
|
|
return nil, ErrLockNotAcquired
|
|
}
|
|
return &Lock{rdb: rdb, key: key, token: token}, nil
|
|
}
|
|
|
|
// Unlock 返回 false 表示锁已经过期或被别人持有（说明业务执行超过了 ttl，需要排查）
|
|
func (l *Lock) Unlock(ctx context.Context) (bool, error) {
|
|
n, err := unlockScript.Run(ctx, l.rdb, []string{l.key}, l.token).Int()
|
|
return n == 1, err
|
|
}
|