modelRT/distributedlock/redis_rwlock.go

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)
}