You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
110 lines
2.4 KiB
Go
110 lines
2.4 KiB
Go
package redis
|
|
|
|
import (
|
|
"context"
|
|
"math/rand"
|
|
"strconv"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
red "github.com/go-redis/redis/v8"
|
|
|
|
"github.com/zeromicro/go-zero/core/logx"
|
|
"github.com/zeromicro/go-zero/core/stringx"
|
|
)
|
|
|
|
const (
|
|
randomLen = 16
|
|
tolerance = 500 // milliseconds
|
|
millisPerSecond = 1000
|
|
)
|
|
|
|
var (
|
|
lockScript = NewScript(`if redis.call("GET", KEYS[1]) == ARGV[1] then
|
|
redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2])
|
|
return "OK"
|
|
else
|
|
return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2])
|
|
end`)
|
|
delScript = NewScript(`if redis.call("GET", KEYS[1]) == ARGV[1] then
|
|
return redis.call("DEL", KEYS[1])
|
|
else
|
|
return 0
|
|
end`)
|
|
)
|
|
|
|
// A RedisLock is a redis lock.
|
|
type RedisLock struct {
|
|
store *Redis
|
|
seconds uint32
|
|
key string
|
|
id string
|
|
}
|
|
|
|
func init() {
|
|
rand.Seed(time.Now().UnixNano())
|
|
}
|
|
|
|
// NewRedisLock returns a RedisLock.
|
|
func NewRedisLock(store *Redis, key string) *RedisLock {
|
|
return &RedisLock{
|
|
store: store,
|
|
key: key,
|
|
id: stringx.Randn(randomLen),
|
|
}
|
|
}
|
|
|
|
// Acquire acquires the lock.
|
|
func (rl *RedisLock) Acquire() (bool, error) {
|
|
return rl.AcquireCtx(context.Background())
|
|
}
|
|
|
|
// AcquireCtx acquires the lock with the given ctx.
|
|
func (rl *RedisLock) AcquireCtx(ctx context.Context) (bool, error) {
|
|
seconds := atomic.LoadUint32(&rl.seconds)
|
|
resp, err := rl.store.ScriptRunCtx(ctx, lockScript, []string{rl.key}, []string{
|
|
rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
|
|
})
|
|
if err == red.Nil {
|
|
return false, nil
|
|
} else if err != nil {
|
|
logx.Errorf("Error on acquiring lock for %s, %s", rl.key, err.Error())
|
|
return false, err
|
|
} else if resp == nil {
|
|
return false, nil
|
|
}
|
|
|
|
reply, ok := resp.(string)
|
|
if ok && reply == "OK" {
|
|
return true, nil
|
|
}
|
|
|
|
logx.Errorf("Unknown reply when acquiring lock for %s: %v", rl.key, resp)
|
|
return false, nil
|
|
}
|
|
|
|
// Release releases the lock.
|
|
func (rl *RedisLock) Release() (bool, error) {
|
|
return rl.ReleaseCtx(context.Background())
|
|
}
|
|
|
|
// ReleaseCtx releases the lock with the given ctx.
|
|
func (rl *RedisLock) ReleaseCtx(ctx context.Context) (bool, error) {
|
|
resp, err := rl.store.ScriptRunCtx(ctx, delScript, []string{rl.key}, []string{rl.id})
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
|
|
reply, ok := resp.(int64)
|
|
if !ok {
|
|
return false, nil
|
|
}
|
|
|
|
return reply == 1, nil
|
|
}
|
|
|
|
// SetExpire sets the expiration.
|
|
func (rl *RedisLock) SetExpire(seconds int) {
|
|
atomic.StoreUint32(&rl.seconds, uint32(seconds))
|
|
}
|