package distributedlock import ( "context" "errors" "fmt" "strings" "sync" "time" constants "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).(*constants.RedisResult) if result.Code == constants.UnknownInternalError { logger.Error(ctx, result.OutputResultMessage()) return fmt.Errorf("get read lock failed:%w", result) } if result.Code == constants.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).(*constants.RedisResult) if (result.Code == constants.RLockFailureWithWLockOccupancy) || (result.Code == constants.UnknownInternalError) { logger.Info(ctx, result.OutputResultMessage()) continue } if result.Code == constants.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 constants.AcquireTimeoutErr } } } return fmt.Errorf("lock the redis read lock failed:%w", result) } func (rl *RedissionRWLocker) tryRLock(ctx context.Context) error { lockType := constants.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 constants.NewRedisResult(constants.UnknownInternalError, lockType, err.Error()) } return constants.NewRedisResult(constants.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 constants.RedisCode(val) == constants.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 constants.RedisCode(val) == constants.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", constants.NewRedisResult(constants.UnknownInternalError, constants.UnRLockType, err.Error())) } if (constants.RedisCode(val) == constants.UnLockSuccess) || (constants.RedisCode(val) == constants.UnRLockSuccess) { if rl.needRefresh && (constants.RedisCode(val) == constants.UnLockSuccess) { rl.cancelRefreshLockTime() } logger.Info(ctx, "unlock read lock success", "token", rl.Token, "key", rl.Key) return nil } if constants.RedisCode(val) == constants.UnRLockFailureWithWLockOccupancy { logger.Info(ctx, "unlock read lock failed", "token", rl.Token, "key", rl.Key) return fmt.Errorf("unlock read lock failed:%w", constants.NewRedisResult(constants.UnRLockFailureWithWLockOccupancy, constants.UnRLockType, "")) } return nil } func (rl *RedissionRWLocker) WLock(ctx context.Context, timeout ...time.Duration) error { result := rl.tryWLock(ctx).(*constants.RedisResult) if result.Code == constants.UnknownInternalError { logger.Error(ctx, result.OutputResultMessage()) return fmt.Errorf("get write lock failed:%w", result) } if result.Code == constants.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).(*constants.RedisResult) if (result.Code == constants.UnknownInternalError) || (result.Code == constants.WLockFailureWithRLockOccupancy) || (result.Code == constants.WLockFailureWithWLockOccupancy) || (result.Code == constants.WLockFailureWithNotFirstPriority) { logger.Info(ctx, result.OutputResultMessage()) continue } if result.Code == constants.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 constants.AcquireTimeoutErr } } } return fmt.Errorf("lock write lock failed:%w", result) } func (rl *RedissionRWLocker) tryWLock(ctx context.Context) error { lockType := constants.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 constants.NewRedisResult(constants.UnknownInternalError, lockType, err.Error()) } return constants.NewRedisResult(constants.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", constants.NewRedisResult(constants.UnknownInternalError, constants.UnWLockType, err.Error())) } if (constants.RedisCode(val) == constants.UnLockSuccess) || constants.RedisCode(val) == constants.UnWLockSuccess { if rl.needRefresh && (constants.RedisCode(val) == constants.UnLockSuccess) { rl.cancelRefreshLockTime() } logger.Info(ctx, "unlock write lock success", "token", rl.Token, "key", rl.Key) return nil } if (constants.RedisCode(val) == constants.UnWLockFailureWithRLockOccupancy) || (constants.RedisCode(val) == constants.UnWLockFailureWithWLockOccupancy) { logger.Info(ctx, "unlock write lock failed", "token", rl.Token, "key", rl.Key) return fmt.Errorf("unlock write lock failed:%w", constants.NewRedisResult(constants.RedisCode(val), constants.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) }