330 lines
11 KiB
Go
330 lines
11 KiB
Go
package distributedlock
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"modelRT/distributedlock/constant"
|
|
"modelRT/distributedlock/luascript"
|
|
"modelRT/logger"
|
|
|
|
uuid "github.com/gofrs/uuid"
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
type RedissionRWLocker struct {
|
|
redissionLocker
|
|
writeWaitChanKey string
|
|
readWaitChanKey string
|
|
RWTokenTimeoutPrefix string
|
|
}
|
|
|
|
func (rl *RedissionRWLocker) RLock(ctx context.Context, timeout ...time.Duration) error {
|
|
result := rl.tryRLock(ctx).(*constant.RedisResult)
|
|
if result.Code == constant.UnknownInternalError {
|
|
logger.Error(ctx, result.OutputResultMessage())
|
|
return fmt.Errorf("get read lock failed:%w", result)
|
|
}
|
|
|
|
if result.Code == constant.LockSuccess {
|
|
if rl.needRefresh {
|
|
rl.refreshOnce.Do(func() {
|
|
if rl.refreshExitChan == nil {
|
|
rl.refreshExitChan = make(chan struct{})
|
|
}
|
|
|
|
// async refresh lock timeout unitl receive exit singal
|
|
go rl.refreshLockTimeout(ctx)
|
|
})
|
|
}
|
|
logger.Info(ctx, "success get the read lock by key and token", "key", rl.Key, "token", rl.Token)
|
|
return nil
|
|
}
|
|
|
|
if len(timeout) > 0 && timeout[0] > 0 {
|
|
if rl.subExitChan == nil {
|
|
rl.subExitChan = make(chan struct{})
|
|
}
|
|
|
|
subMsgChan := make(chan struct{}, 1)
|
|
sub := rl.client.Subscribe(ctx, rl.readWaitChanKey)
|
|
go rl.subscribeLock(ctx, sub, subMsgChan)
|
|
|
|
acquireTimer := time.NewTimer(timeout[0])
|
|
for {
|
|
select {
|
|
case _, ok := <-subMsgChan:
|
|
if !ok {
|
|
err := errors.New("failed to read the read lock waiting for for the channel message")
|
|
logger.Error(ctx, "failed to read the read lock waiting for for the channel message")
|
|
return err
|
|
}
|
|
|
|
result := rl.tryRLock(ctx).(*constant.RedisResult)
|
|
if (result.Code == constant.RLockFailureWithWLockOccupancy) || (result.Code == constant.UnknownInternalError) {
|
|
logger.Info(ctx, result.OutputResultMessage())
|
|
continue
|
|
}
|
|
|
|
if result.Code == constant.LockSuccess {
|
|
logger.Info(ctx, result.OutputResultMessage())
|
|
rl.closeSub(ctx, sub, rl.subExitChan)
|
|
|
|
if rl.needRefresh {
|
|
rl.refreshOnce.Do(func() {
|
|
if rl.refreshExitChan == nil {
|
|
rl.refreshExitChan = make(chan struct{})
|
|
}
|
|
|
|
// async refresh lock timeout unitl receive exit singal
|
|
go rl.refreshLockTimeout(ctx)
|
|
})
|
|
}
|
|
return nil
|
|
}
|
|
case <-acquireTimer.C:
|
|
logger.Info(ctx, "the waiting time for obtaining the read lock operation has timed out")
|
|
rl.closeSub(ctx, sub, rl.subExitChan)
|
|
// after acquire lock timeout,notice the sub channel to close
|
|
return constant.AcquireTimeoutErr
|
|
}
|
|
}
|
|
}
|
|
return fmt.Errorf("lock the redis read lock failed:%w", result)
|
|
}
|
|
|
|
func (rl *RedissionRWLocker) tryRLock(ctx context.Context) error {
|
|
lockType := constant.LockType
|
|
|
|
res := rl.client.Eval(ctx, luascript.RLockScript, []string{rl.Key, rl.RWTokenTimeoutPrefix}, rl.lockLeaseTime, rl.Token)
|
|
val, err := res.Int()
|
|
if err != redis.Nil && err != nil {
|
|
return constant.NewRedisResult(constant.UnknownInternalError, lockType, err.Error())
|
|
}
|
|
return constant.NewRedisResult(constant.RedisCode(val), lockType, "")
|
|
}
|
|
|
|
func (rl *RedissionRWLocker) refreshLockTimeout(ctx context.Context) {
|
|
logger.Info(ctx, "lock refresh by key and token", "token", rl.Token, "key", rl.Key)
|
|
|
|
lockTime := time.Duration(rl.lockLeaseTime/3) * time.Millisecond
|
|
timer := time.NewTimer(lockTime)
|
|
defer timer.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-timer.C:
|
|
// extend key lease time
|
|
res := rl.client.Eval(ctx, luascript.RefreshRWLockScript, []string{rl.Key, rl.RWTokenTimeoutPrefix}, rl.lockLeaseTime, rl.Token)
|
|
val, err := res.Int()
|
|
if err != redis.Nil && err != nil {
|
|
logger.Info(ctx, "lock refresh failed", "token", rl.Token, "key", rl.Key, "error", err)
|
|
return
|
|
}
|
|
|
|
if constant.RedisCode(val) == constant.RefreshLockFailure {
|
|
logger.Error(ctx, "lock refreash failed,can not find the read lock by key and token", "rwTokenPrefix", rl.RWTokenTimeoutPrefix, "token", rl.Token, "key", rl.Key)
|
|
return
|
|
}
|
|
|
|
if constant.RedisCode(val) == constant.RefreshLockSuccess {
|
|
logger.Info(ctx, "lock refresh success by key and token", "token", rl.Token, "key", rl.Key)
|
|
}
|
|
timer.Reset(lockTime)
|
|
case <-rl.refreshExitChan:
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (rl *RedissionRWLocker) UnRLock(ctx context.Context) error {
|
|
logger.Info(ctx, "unlock RLock by key and token", "key", rl.Key, "token", rl.Token)
|
|
res := rl.client.Eval(ctx, luascript.UnRLockScript, []string{rl.Key, rl.RWTokenTimeoutPrefix, rl.writeWaitChanKey}, unlockMessage, rl.Token)
|
|
val, err := res.Int()
|
|
if err != redis.Nil && err != nil {
|
|
logger.Info(ctx, "unlock read lock failed", "token", rl.Token, "key", rl.Key, "error", err)
|
|
return fmt.Errorf("unlock read lock failed:%w", constant.NewRedisResult(constant.UnknownInternalError, constant.UnRLockType, err.Error()))
|
|
}
|
|
|
|
if (constant.RedisCode(val) == constant.UnLockSuccess) || (constant.RedisCode(val) == constant.UnRLockSuccess) {
|
|
if rl.needRefresh && (constant.RedisCode(val) == constant.UnLockSuccess) {
|
|
rl.cancelRefreshLockTime()
|
|
}
|
|
|
|
logger.Info(ctx, "unlock read lock success", "token", rl.Token, "key", rl.Key)
|
|
return nil
|
|
}
|
|
|
|
if constant.RedisCode(val) == constant.UnRLockFailureWithWLockOccupancy {
|
|
logger.Info(ctx, "unlock read lock failed", "token", rl.Token, "key", rl.Key)
|
|
return fmt.Errorf("unlock read lock failed:%w", constant.NewRedisResult(constant.UnRLockFailureWithWLockOccupancy, constant.UnRLockType, ""))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (rl *RedissionRWLocker) WLock(ctx context.Context, timeout ...time.Duration) error {
|
|
result := rl.tryWLock(ctx).(*constant.RedisResult)
|
|
if result.Code == constant.UnknownInternalError {
|
|
logger.Error(ctx, result.OutputResultMessage())
|
|
return fmt.Errorf("get write lock failed:%w", result)
|
|
}
|
|
|
|
if result.Code == constant.LockSuccess {
|
|
if rl.needRefresh {
|
|
rl.refreshOnce.Do(func() {
|
|
if rl.refreshExitChan == nil {
|
|
rl.refreshExitChan = make(chan struct{})
|
|
}
|
|
|
|
// async refresh lock timeout unitl receive exit singal
|
|
go rl.refreshLockTimeout(ctx)
|
|
})
|
|
}
|
|
logger.Info(ctx, "success get the write lock by key and token", "key", rl.Key, "token", rl.Token)
|
|
return nil
|
|
}
|
|
|
|
if len(timeout) > 0 && timeout[0] > 0 {
|
|
if rl.subExitChan == nil {
|
|
rl.subExitChan = make(chan struct{})
|
|
}
|
|
|
|
subMsgChan := make(chan struct{}, 1)
|
|
sub := rl.client.Subscribe(ctx, rl.writeWaitChanKey)
|
|
go rl.subscribeLock(ctx, sub, subMsgChan)
|
|
|
|
acquireTimer := time.NewTimer(timeout[0])
|
|
for {
|
|
select {
|
|
case _, ok := <-subMsgChan:
|
|
if !ok {
|
|
err := errors.New("failed to read the write lock waiting for for the channel message")
|
|
logger.Error(ctx, "failed to read the read lock waiting for for the channel message")
|
|
return err
|
|
}
|
|
|
|
result := rl.tryWLock(ctx).(*constant.RedisResult)
|
|
if (result.Code == constant.UnknownInternalError) || (result.Code == constant.WLockFailureWithRLockOccupancy) || (result.Code == constant.WLockFailureWithWLockOccupancy) || (result.Code == constant.WLockFailureWithNotFirstPriority) {
|
|
logger.Info(ctx, result.OutputResultMessage())
|
|
continue
|
|
}
|
|
|
|
if result.Code == constant.LockSuccess {
|
|
logger.Info(ctx, result.OutputResultMessage())
|
|
rl.closeSub(ctx, sub, rl.subExitChan)
|
|
|
|
if rl.needRefresh {
|
|
rl.refreshOnce.Do(func() {
|
|
if rl.refreshExitChan == nil {
|
|
rl.refreshExitChan = make(chan struct{})
|
|
}
|
|
|
|
// async refresh lock timeout unitl receive exit singal
|
|
go rl.refreshLockTimeout(ctx)
|
|
})
|
|
}
|
|
return nil
|
|
}
|
|
case <-acquireTimer.C:
|
|
logger.Info(ctx, "the waiting time for obtaining the write lock operation has timed out")
|
|
rl.closeSub(ctx, sub, rl.subExitChan)
|
|
// after acquire lock timeout,notice the sub channel to close
|
|
return constant.AcquireTimeoutErr
|
|
}
|
|
}
|
|
}
|
|
return fmt.Errorf("lock write lock failed:%w", result)
|
|
}
|
|
|
|
func (rl *RedissionRWLocker) tryWLock(ctx context.Context) error {
|
|
lockType := constant.LockType
|
|
|
|
res := rl.client.Eval(ctx, luascript.WLockScript, []string{rl.Key, rl.RWTokenTimeoutPrefix}, rl.lockLeaseTime, rl.Token)
|
|
val, err := res.Int()
|
|
if err != redis.Nil && err != nil {
|
|
return constant.NewRedisResult(constant.UnknownInternalError, lockType, err.Error())
|
|
}
|
|
return constant.NewRedisResult(constant.RedisCode(val), lockType, "")
|
|
}
|
|
|
|
func (rl *RedissionRWLocker) UnWLock(ctx context.Context) error {
|
|
res := rl.client.Eval(ctx, luascript.UnWLockScript, []string{rl.Key, rl.RWTokenTimeoutPrefix, rl.writeWaitChanKey, rl.readWaitChanKey}, unlockMessage, rl.Token)
|
|
val, err := res.Int()
|
|
if err != redis.Nil && err != nil {
|
|
logger.Error(ctx, "unlock write lock failed", "token", rl.Token, "key", rl.Key, "error", err)
|
|
return fmt.Errorf("unlock write lock failed:%w", constant.NewRedisResult(constant.UnknownInternalError, constant.UnWLockType, err.Error()))
|
|
}
|
|
|
|
if (constant.RedisCode(val) == constant.UnLockSuccess) || constant.RedisCode(val) == constant.UnWLockSuccess {
|
|
if rl.needRefresh && (constant.RedisCode(val) == constant.UnLockSuccess) {
|
|
rl.cancelRefreshLockTime()
|
|
}
|
|
logger.Info(ctx, "unlock write lock success", "token", rl.Token, "key", rl.Key)
|
|
return nil
|
|
}
|
|
|
|
if (constant.RedisCode(val) == constant.UnWLockFailureWithRLockOccupancy) || (constant.RedisCode(val) == constant.UnWLockFailureWithWLockOccupancy) {
|
|
logger.Info(ctx, "unlock write lock failed", "token", rl.Token, "key", rl.Key)
|
|
return fmt.Errorf("unlock write lock failed:%w", constant.NewRedisResult(constant.RedisCode(val), constant.UnWLockType, ""))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// TODO 优化 panic
|
|
func GetRWLocker(client *redis.Client, conf *RedissionLockConfig) *RedissionRWLocker {
|
|
if conf.Token == "" {
|
|
token, err := uuid.NewV4()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
conf.Token = token.String()
|
|
}
|
|
|
|
if conf.Prefix == "" {
|
|
conf.Prefix = "redission-rwlock"
|
|
}
|
|
|
|
if conf.TimeoutPrefix == "" {
|
|
conf.TimeoutPrefix = "rwlock_timeout"
|
|
}
|
|
|
|
if conf.ChanPrefix == "" {
|
|
conf.ChanPrefix = "redission-rwlock-channel"
|
|
}
|
|
|
|
if conf.LockLeaseTime == 0 {
|
|
conf.LockLeaseTime = internalLockLeaseTime
|
|
}
|
|
|
|
r := &redissionLocker{
|
|
Token: conf.Token,
|
|
Key: strings.Join([]string{conf.Prefix, conf.Key}, ":"),
|
|
needRefresh: conf.NeedRefresh,
|
|
lockLeaseTime: conf.LockLeaseTime,
|
|
client: client,
|
|
refreshOnce: &sync.Once{},
|
|
}
|
|
|
|
rwLocker := &RedissionRWLocker{
|
|
redissionLocker: *r,
|
|
writeWaitChanKey: strings.Join([]string{conf.ChanPrefix, conf.Key, "write"}, ":"),
|
|
readWaitChanKey: strings.Join([]string{conf.ChanPrefix, conf.Key, "read"}, ":"),
|
|
RWTokenTimeoutPrefix: conf.TimeoutPrefix,
|
|
}
|
|
return rwLocker
|
|
}
|
|
|
|
func InitRWLocker(key string, token string, lockLeaseTime uint64, needRefresh bool) *RedissionRWLocker {
|
|
conf := &RedissionLockConfig{
|
|
Key: key,
|
|
Token: token,
|
|
LockLeaseTime: lockLeaseTime,
|
|
NeedRefresh: needRefresh,
|
|
}
|
|
return GetRWLocker(GetRedisClientInstance(), conf)
|
|
}
|