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 { if rl.exit == nil { rl.exit = make(chan struct{}) } 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.once.Do(func() { // 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 } subMsg := make(chan struct{}, 1) defer close(subMsg) sub := rl.client.Subscribe(ctx, rl.readWaitChanKey) defer sub.Close() go rl.subscribeLock(ctx, sub, subMsg) if len(timeout) > 0 && timeout[0] > 0 { acquireTimer := time.NewTimer(timeout[0]) for { select { case _, ok := <-subMsg: 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 } resultErr := rl.tryRLock(ctx).(*constant.RedisResult) if (resultErr.Code == constant.RLockFailureWithWLockOccupancy) || (resultErr.Code == constant.UnknownInternalError) { rl.logger.Info(resultErr.OutputResultMessage()) continue } if resultErr.Code == constant.LockSuccess { rl.logger.Info(resultErr.OutputResultMessage()) return nil } case <-acquireTimer.C: err := errors.New("the waiting time for obtaining the read lock operation has timed out") rl.logger.Info("the waiting time for obtaining the read lock operation has timed out") return err } } } 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.exit: 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 { if rl.exit == nil { rl.exit = make(chan struct{}) } 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.once.Do(func() { // 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 } subMsg := make(chan struct{}, 1) defer close(subMsg) sub := rl.client.Subscribe(ctx, rl.writeWaitChanKey) defer sub.Close() go rl.subscribeLock(ctx, sub, subMsg) if len(timeout) > 0 && timeout[0] > 0 { acquireTimer := time.NewTimer(timeout[0]) for { select { case _, ok := <-subMsg: 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()) return nil } case <-acquireTimer.C: rl.logger.Info("the waiting time for obtaining the write lock operation has timed out") 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, exit: make(chan struct{}), once: &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) }