utils/redislock/lock.go

package redislock

import (
	"context"
	"github.com/pkg/errors"
	"github.com/redis/go-redis/v9"
	"time"
)

const (
	tryLockTimeoutSec = 5
)

type Lock struct {
	redisClient *redis.Client
	ownerID     string
}

func NewLock(address string, password string, db int) (*Lock, error) {
	if address == "" {
		return nil, errors.New("redis address不能为空")
	}

	return &Lock{redisClient: redis.NewClient(&redis.Options{
		Addr:     address,
		Password: password,
		DB:       db,
	})}, nil
}

func DestroyLock(lock *Lock) error {
	if lock == nil || lock.redisClient == nil {
		return nil
	}

	err := lock.redisClient.Close()
	if err != nil {
		return err
	}

	lock.redisClient = nil

	return nil
}

func (lock *Lock) Lock(ownerID string, key string, expireSec int) error {
	if ownerID == "" {
		return errors.New("ownerID不能为空")
	}

	if key == "" {
		return errors.New("key不能为空")
	}

	if expireSec == 0 {
		return errors.New("expireSec不能为0")
	}

	for {
		locked, err := lock.tryLock(ownerID, key, expireSec)
		if err != nil {
			return err
		}

		if locked {
			return nil
		}

		time.Sleep(200 * time.Millisecond)
	}
}

func (lock *Lock) TryLock(ownerID string, key string, expireSec int) (bool, error) {
	if ownerID == "" {
		return false, errors.New("ownerID不能为空")
	}

	if key == "" {
		return false, errors.New("key不能为空")
	}

	if expireSec == 0 {
		return false, errors.New("expireSec不能为0")
	}

	return lock.tryLock(ownerID, key, expireSec)
}

func (lock *Lock) Unlock(ownerID string, key string) error {
	if ownerID == "" {
		return errors.New("ownerID不能为空")
	}

	if key == "" {
		return errors.New("key不能为空")
	}

	if ownerID != lock.ownerID {
		return nil
	}

	return lock.deleteLockTime(key)
}

func (lock *Lock) tryLock(ownerID string, key string, expireSec int) (bool, error) {
	expire := time.Duration(expireSec) * time.Second
	nowTime := time.Now()
	nowNano := nowTime.UnixNano()
	currentLockNano := nowTime.Add(expire).UnixNano()

	// 尝试以当前lockTime上锁，如果返回true，则上锁成功，否则锁正被占用
	locked, err := lock.setLockTime(key, currentLockNano, expire)
	if err != nil {
		return false, err
	}

	if locked {
		lock.ownerID = ownerID
		return true, nil
	}

	// 获取保存的lockTime
	savedLockTimeNano, err := lock.getLockTime(key)
	if err != nil {
		return false, err
	}

	// 锁已经过期
	if savedLockTimeNano != 0 && savedLockTimeNano < nowNano {
		// 尝试以当前lockTime上锁，返回之前设置的值
		oldLockTime, err := lock.getAndSetLockTime(key, currentLockNano)
		if err != nil {
			return false, err
		}

		// 排除上一步操作存在的并发问题
		if oldLockTime != 0 && oldLockTime == savedLockTimeNano {
			lock.ownerID = ownerID
			return true, nil
		}
	}

	return false, nil
}

func (lock *Lock) setLockTime(key string, lockTime int64, expireSec time.Duration) (bool, error) {
	ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
	defer cancel()

	cmd := lock.redisClient.SetNX(ctx, key, lockTime, expireSec)
	if cmd.Err() != nil {
		return false, cmd.Err()
	}

	return cmd.Val(), nil
}

func (lock *Lock) getLockTime(key string) (int64, error) {
	ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
	defer cancel()

	cmd := lock.redisClient.Get(ctx, key)
	if cmd.Err() != nil {
		if cmd.Err().Error() == redis.Nil.Error() {
			return 0, nil
		}

		return 0, cmd.Err()
	}

	lockTime, err := cmd.Int64()
	if err != nil {
		return 0, err
	}

	return lockTime, nil
}

func (lock *Lock) getAndSetLockTime(key string, newLockTime int64) (int64, error) {
	ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
	defer cancel()

	cmd := lock.redisClient.GetSet(ctx, key, newLockTime)
	if cmd.Err() != nil {
		if cmd.Err().Error() == redis.Nil.Error() {
			return 0, nil
		}

		return 0, cmd.Err()
	}

	lockTime, err := cmd.Int64()
	if err != nil {
		return 0, err
	}

	return lockTime, nil
}

func (lock *Lock) deleteLockTime(key string) error {
	ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
	defer cancel()

	cmd := lock.redisClient.Del(ctx, key)
	if cmd.Err() != nil {
		return cmd.Err()
	}

	return nil
}