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" "go.uber.org/zap" ) 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 { rl.logger.Error(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) }) } rl.logger.Info("success get the read lock by key and token", zap.String("key", rl.key), zap.String("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(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") rl.logger.Error("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) { rl.logger.Info(result.OutputResultMessage()) continue } if result.Code == constant.LockSuccess { rl.logger.Info(result.OutputResultMessage()) rl.closeSub(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: rl.logger.Info("the waiting time for obtaining the read lock operation has timed out") rl.closeSub(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) { rl.logger.Info("lock refresh by key and token", zap.String("token", rl.token), zap.String("key", rl.key)) lockTime := time.Duration(rl.lockLeaseTime/3) * time.Second 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 { rl.logger.Info("lock refresh failed", zap.String("token", rl.token), zap.String("key", rl.key), zap.Error(err)) return } if constant.RedisCode(val) == constant.RefreshLockFailure { rl.logger.Error("lock refreash failed,can not find the read lock by key and token", zap.String("rwTokenPrefix", rl.rwTokenTimeoutPrefix), zap.String("token", rl.token), zap.String("key", rl.key)) return } if constant.RedisCode(val) == constant.RefreshLockSuccess { rl.logger.Info("lock refresh success by key and token", zap.String("token", rl.token), zap.String("key", rl.key)) } timer.Reset(lockTime) case <-rl.refreshExitChan: return } } } func (rl *RedissionRWLocker) UnRLock(ctx context.Context) error { rl.logger.Info("unlock RLock by key and token", zap.String("key", rl.key), zap.String("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 { rl.logger.Info("unlock read lock failed", zap.String("token", rl.token), zap.String("key", rl.key), zap.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() } rl.logger.Info("unlock read lock success", zap.String("token", rl.token), zap.String("key", rl.key)) return nil } if constant.RedisCode(val) == constant.UnRLockFailureWithWLockOccupancy { rl.logger.Info("unlock read lock failed", zap.String("token", rl.token), zap.String("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 { rl.logger.Error(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) }) } rl.logger.Info("success get the write lock by key and token", zap.String("key", rl.key), zap.String("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(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") rl.logger.Error("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) { rl.logger.Info(result.OutputResultMessage()) continue } if result.Code == constant.LockSuccess { rl.logger.Info(result.OutputResultMessage()) rl.closeSub(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: rl.logger.Info("the waiting time for obtaining the write lock operation has timed out") rl.closeSub(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 { rl.logger.Error("unlock write lock failed", zap.String("token", rl.token), zap.String("key", rl.key), zap.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() } rl.logger.Info("unlock write lock success", zap.String("token", rl.token), zap.String("key", rl.key)) return nil } if (constant.RedisCode(val) == constant.UnWLockFailureWithRLockOccupancy) || (constant.RedisCode(val) == constant.UnWLockFailureWithWLockOccupancy) { rl.logger.Info("unlock write lock failed", zap.String("token", rl.token), zap.String("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{}, logger: logger.GetLoggerInstance(), } 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) }