-
-
Notifications
You must be signed in to change notification settings - Fork 3.8k
/
redislock.go
95 lines (83 loc) · 2 KB
/
redislock.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
package redis
import (
"math/rand"
"strconv"
"sync/atomic"
"time"
red "github.com/go-redis/redis"
"github.com/tal-tech/go-zero/core/logx"
)
const (
letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
lockCommand = `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`
delCommand = `if redis.call("GET", KEYS[1]) == ARGV[1] then
return redis.call("DEL", KEYS[1])
else
return 0
end`
randomLen = 16
tolerance = 500 // milliseconds
millisPerSecond = 1000
)
type RedisLock struct {
store *Redis
seconds uint32
key string
id string
}
func init() {
rand.Seed(time.Now().UnixNano())
}
func NewRedisLock(store *Redis, key string) *RedisLock {
return &RedisLock{
store: store,
key: key,
id: randomStr(randomLen),
}
}
func (rl *RedisLock) Acquire() (bool, error) {
seconds := atomic.LoadUint32(&rl.seconds)
resp, err := rl.store.Eval(lockCommand, []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
} else {
logx.Errorf("Unknown reply when acquiring lock for %s: %v", rl.key, resp)
return false, nil
}
}
func (rl *RedisLock) Release() (bool, error) {
resp, err := rl.store.Eval(delCommand, []string{rl.key}, []string{rl.id})
if err != nil {
return false, err
}
if reply, ok := resp.(int64); !ok {
return false, nil
} else {
return reply == 1, nil
}
}
func (rl *RedisLock) SetExpire(seconds int) {
atomic.StoreUint32(&rl.seconds, uint32(seconds))
}
func randomStr(n int) string {
b := make([]byte, n)
for i := range b {
b[i] = letters[rand.Intn(len(letters))]
}
return string(b)
}