modelRT/distributedlock/redis_rwlock.go

package distributed_lock

import (
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"modelRT/distributedlock/constant"
	"modelRT/distributedlock/luascript"
	"modelRT/logger"

	"github.com/go-redis/redis"
	uuid "github.com/google/uuid"
	"go.uber.org/zap"
)

type RedissionRWLocker struct {
	redissionLocker
	rwTokenTimeoutPrefix string
}

func (rl *RedissionRWLocker) RLock(timeout ...time.Duration) error {
	if rl.exit == nil {
		rl.exit = make(chan struct{})
	}

	result := rl.tryRLock().(*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()
			})
		}
		rl.logger.Info("success get the read 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(rl.waitChanKey)
	defer sub.Close()
	go rl.subscribeLock(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().(*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() error {
	lockType := constant.LockType

	res := rl.client.Eval(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() {
	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(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() error {
	rl.logger.Info("unlock RLock by key and token", zap.String("key", rl.key), zap.String("token", rl.token))
	res := rl.client.Eval(luascript.UnRLockScript, []string{rl.key, rl.rwTokenTimeoutPrefix, rl.waitChanKey}, 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.UnLockType, err.Error()))
	}

	if (constant.RedisCode(val) == constant.UnLockSuccess) || (constant.RedisCode(val) == constant.UnRLockSuccess) {
		if rl.needRefresh {
			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.UnLockType, ""))
	}
	return nil
}

func (rl *RedissionRWLocker) WLock(timeout ...time.Duration) error {
	if rl.exit == nil {
		rl.exit = make(chan struct{})
	}

	result := rl.tryWLock().(*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) && rl.needRefresh {
		rl.once.Do(func() {
			// async refresh lock timeout unitl receive exit singal
			go rl.refreshLockTimeout()
		})
		return nil
	}

	subMsg := make(chan struct{}, 1)
	defer close(subMsg)
	sub := rl.client.Subscribe(rl.waitChanKey)
	defer sub.Close()
	go rl.subscribeLock(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().(*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:
				err := errors.New("the waiting time for obtaining the write lock operation has timed out")
				rl.logger.Info("the waiting time for obtaining the write lock operation has timed out")
				return err
			}
		}
	}
	return fmt.Errorf("lock write lock failed:%w", result)
}

func (rl *RedissionRWLocker) tryWLock() error {
	lockType := constant.LockType

	res := rl.client.Eval(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() error {
	res := rl.client.Eval(luascript.UnWLockScript, []string{rl.key, rl.rwTokenTimeoutPrefix, rl.waitChanKey}, unlockMessage, rl.token)
	val, err := res.Int()
	if err != redis.Nil && err != nil {
		rl.logger.Info("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.UnLockType, err.Error()))
	}

	if constant.RedisCode(val) == constant.UnLockSuccess {
		if rl.needRefresh {
			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.UnLockType, ""))
	}
	return nil
}

func GetRWLocker(client *redis.Client, ops *RedissionLockConfig) *RedissionRWLocker {
	if ops.Token == "" {
		ops.Token = uuid.New().String()
	}

	if ops.Prefix == "" {
		ops.Prefix = "redission-rwlock"
	}

	if ops.TimeoutPrefix == "" {
		ops.TimeoutPrefix = "rwlock_timeout"
	}

	if ops.ChanPrefix == "" {
		ops.ChanPrefix = "redission-rwlock-channel"
	}

	if ops.LockLeaseTime == 0 {
		ops.LockLeaseTime = internalLockLeaseTime
	}

	r := &redissionLocker{
		token:         ops.Token,
		key:           strings.Join([]string{ops.Prefix, ops.Key}, ":"),
		needRefresh:   ops.NeedRefresh,
		lockLeaseTime: ops.LockLeaseTime,
		waitChanKey:   strings.Join([]string{ops.ChanPrefix, ops.Key, "write"}, ":"),
		client:        client,
		exit:          make(chan struct{}),
		once:          &sync.Once{},
		logger:        logger.GetLoggerInstance(),
	}

	rwLocker := &RedissionRWLocker{
		redissionLocker:      *r,
		rwTokenTimeoutPrefix: ops.TimeoutPrefix,
	}
	return rwLocker
}