2023-06-25 07:58:21 +00:00
|
|
|
package lock_manager
|
|
|
|
|
|
|
|
import (
|
|
|
|
"fmt"
|
|
|
|
"github.com/google/uuid"
|
|
|
|
"github.com/puzpuzpuz/xsync/v2"
|
|
|
|
"time"
|
|
|
|
)
|
|
|
|
|
2023-06-26 00:37:54 +00:00
|
|
|
var LockErrorNonEmptyTokenOnNewLock = fmt.Errorf("lock: non-empty token on a new lock")
|
|
|
|
var LockErrorNonEmptyTokenOnExpiredLock = fmt.Errorf("lock: non-empty token on an expired lock")
|
|
|
|
var LockErrorTokenMismatch = fmt.Errorf("lock: token mismatch")
|
|
|
|
var UnlockErrorTokenMismatch = fmt.Errorf("unlock: token mismatch")
|
|
|
|
|
2023-09-16 22:06:16 +00:00
|
|
|
// LockManager local lock manager, used by distributed lock manager
|
2023-06-25 07:58:21 +00:00
|
|
|
type LockManager struct {
|
|
|
|
locks *xsync.MapOf[string, *Lock]
|
|
|
|
}
|
|
|
|
type Lock struct {
|
2023-06-25 21:14:40 +00:00
|
|
|
Token string
|
|
|
|
ExpiredAtNs int64
|
|
|
|
Key string // only used for moving locks
|
2023-09-16 22:06:16 +00:00
|
|
|
Owner string
|
2023-06-25 07:58:21 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
func NewLockManager() *LockManager {
|
|
|
|
t := &LockManager{
|
|
|
|
locks: xsync.NewMapOf[*Lock](),
|
|
|
|
}
|
|
|
|
go t.CleanUp()
|
|
|
|
return t
|
|
|
|
}
|
|
|
|
|
2023-09-16 22:06:16 +00:00
|
|
|
func (lm *LockManager) Lock(path string, expiredAtNs int64, token string, owner string) (renewToken string, err error) {
|
2023-06-25 07:58:21 +00:00
|
|
|
lm.locks.Compute(path, func(oldValue *Lock, loaded bool) (newValue *Lock, delete bool) {
|
|
|
|
if oldValue != nil {
|
2023-06-25 21:14:40 +00:00
|
|
|
if oldValue.ExpiredAtNs > 0 && oldValue.ExpiredAtNs < time.Now().UnixNano() {
|
2023-06-25 07:58:21 +00:00
|
|
|
// lock is expired, set to a new lock
|
2023-06-25 21:14:40 +00:00
|
|
|
if token != "" {
|
2023-06-26 00:37:54 +00:00
|
|
|
err = LockErrorNonEmptyTokenOnExpiredLock
|
2023-06-25 21:14:40 +00:00
|
|
|
return nil, false
|
|
|
|
} else {
|
|
|
|
// new lock
|
|
|
|
renewToken = uuid.New().String()
|
2023-09-16 22:06:16 +00:00
|
|
|
return &Lock{Token: renewToken, ExpiredAtNs: expiredAtNs, Owner: owner}, false
|
2023-06-25 21:14:40 +00:00
|
|
|
}
|
2023-06-25 07:58:21 +00:00
|
|
|
}
|
2023-06-25 21:14:40 +00:00
|
|
|
// not expired
|
2023-06-25 07:58:21 +00:00
|
|
|
if oldValue.Token == token {
|
2023-06-25 21:14:40 +00:00
|
|
|
// token matches, renew the lock
|
2023-06-25 22:28:01 +00:00
|
|
|
renewToken = uuid.New().String()
|
2023-09-16 22:06:16 +00:00
|
|
|
return &Lock{Token: renewToken, ExpiredAtNs: expiredAtNs, Owner: owner}, false
|
2023-06-25 07:58:21 +00:00
|
|
|
} else {
|
2023-06-26 00:37:54 +00:00
|
|
|
err = LockErrorTokenMismatch
|
2023-06-25 07:58:21 +00:00
|
|
|
return oldValue, false
|
|
|
|
}
|
|
|
|
} else {
|
|
|
|
if token == "" {
|
2023-06-25 21:14:40 +00:00
|
|
|
// new lock
|
2023-06-25 07:58:21 +00:00
|
|
|
renewToken = uuid.New().String()
|
2023-09-16 22:06:16 +00:00
|
|
|
return &Lock{Token: renewToken, ExpiredAtNs: expiredAtNs, Owner: owner}, false
|
2023-06-25 07:58:21 +00:00
|
|
|
} else {
|
2023-06-26 00:37:54 +00:00
|
|
|
err = LockErrorNonEmptyTokenOnNewLock
|
2023-06-25 07:58:21 +00:00
|
|
|
return nil, false
|
|
|
|
}
|
|
|
|
}
|
|
|
|
})
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
func (lm *LockManager) Unlock(path string, token string) (isUnlocked bool, err error) {
|
|
|
|
lm.locks.Compute(path, func(oldValue *Lock, loaded bool) (newValue *Lock, delete bool) {
|
|
|
|
if oldValue != nil {
|
|
|
|
now := time.Now()
|
2023-06-25 21:14:40 +00:00
|
|
|
if oldValue.ExpiredAtNs > 0 && oldValue.ExpiredAtNs < now.UnixNano() {
|
2023-06-25 07:58:21 +00:00
|
|
|
// lock is expired, delete it
|
|
|
|
isUnlocked = true
|
|
|
|
return nil, true
|
|
|
|
}
|
|
|
|
if oldValue.Token == token {
|
2023-06-25 21:14:40 +00:00
|
|
|
if oldValue.ExpiredAtNs <= now.UnixNano() {
|
2023-06-25 07:58:21 +00:00
|
|
|
isUnlocked = true
|
|
|
|
return nil, true
|
|
|
|
}
|
|
|
|
return oldValue, false
|
|
|
|
} else {
|
|
|
|
isUnlocked = false
|
2023-06-26 00:37:54 +00:00
|
|
|
err = UnlockErrorTokenMismatch
|
2023-06-25 07:58:21 +00:00
|
|
|
return oldValue, false
|
|
|
|
}
|
|
|
|
} else {
|
|
|
|
isUnlocked = true
|
|
|
|
return nil, true
|
|
|
|
}
|
|
|
|
})
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
func (lm *LockManager) CleanUp() {
|
|
|
|
for {
|
|
|
|
time.Sleep(1 * time.Minute)
|
|
|
|
now := time.Now().UnixNano()
|
|
|
|
lm.locks.Range(func(key string, value *Lock) bool {
|
2023-06-25 22:28:01 +00:00
|
|
|
if value == nil {
|
|
|
|
return true
|
|
|
|
}
|
2023-06-25 21:14:40 +00:00
|
|
|
if now > value.ExpiredAtNs {
|
2023-06-25 07:58:21 +00:00
|
|
|
lm.locks.Delete(key)
|
|
|
|
return true
|
|
|
|
}
|
|
|
|
return true
|
|
|
|
})
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-06-25 21:14:40 +00:00
|
|
|
// SelectLocks takes out locks by key
|
2023-06-25 07:58:21 +00:00
|
|
|
// if keyFn return true, the lock will be taken out
|
2023-06-25 21:14:40 +00:00
|
|
|
func (lm *LockManager) SelectLocks(selectFn func(key string) bool) (locks []*Lock) {
|
2023-06-25 07:58:21 +00:00
|
|
|
now := time.Now().UnixNano()
|
|
|
|
lm.locks.Range(func(key string, lock *Lock) bool {
|
2023-06-25 21:14:40 +00:00
|
|
|
if now > lock.ExpiredAtNs {
|
2023-06-25 07:58:21 +00:00
|
|
|
lm.locks.Delete(key)
|
|
|
|
return true
|
|
|
|
}
|
2023-06-25 21:14:40 +00:00
|
|
|
if selectFn(key) {
|
2023-06-25 07:58:21 +00:00
|
|
|
lm.locks.Delete(key)
|
|
|
|
lock.Key = key
|
|
|
|
locks = append(locks, lock)
|
|
|
|
}
|
|
|
|
return true
|
|
|
|
})
|
|
|
|
return
|
|
|
|
}
|
2023-06-25 21:14:40 +00:00
|
|
|
|
|
|
|
// InsertLock inserts a lock unconditionally
|
2023-09-16 22:06:16 +00:00
|
|
|
func (lm *LockManager) InsertLock(path string, expiredAtNs int64, token string, owner string) {
|
|
|
|
lm.locks.Store(path, &Lock{Token: token, ExpiredAtNs: expiredAtNs, Owner: owner})
|
|
|
|
}
|
|
|
|
|
|
|
|
func (lm *LockManager) GetLockOwner(key string) (owner string, err error) {
|
|
|
|
lm.locks.Range(func(k string, lock *Lock) bool {
|
|
|
|
if k == key {
|
|
|
|
owner = lock.Owner
|
|
|
|
return false
|
|
|
|
}
|
|
|
|
return true
|
|
|
|
})
|
|
|
|
return
|
|
|
|
|
2023-06-25 21:14:40 +00:00
|
|
|
}
|