goRedisDLM/redisLock/lock.go

package redisLock

import (
	"context"
	"errors"
	"fmt"
	"goRedisDLM/utils"
	"log"
	"sync/atomic"
	"time"

	"github.com/go-redis/redis/v8"
)

var ErrorNotMyLock = errors.New("not myLock")

const REDISLOCKPREFIXKEY = "goRedisLock"

func isErrorNotMyLock(err error) bool {
	return errors.Is(err, ErrorNotMyLock)
}

type RedisLockNew struct {
	RedisLockOptions
	client *redis.Client
	ctx    context.Context
	key    string
	token  string

	DogRunning int32
	stopDog    context.CancelFunc
}

func (r *RedisLockNew) IsHeldByCurrent() bool {
	return r.client.Exists(r.ctx, r.getLockKey()).Val() == 1
}

func (r *RedisLockNew) getLockKey() string {
	return REDISLOCKPREFIXKEY + "_" + r.key
}

func (r *RedisLockNew) blockingLock() error {
	// maxWaitSecond为0时认为永久阻塞
	var timeAfter <-chan time.Time
	if r.maxWaitSecond > 0 {
		timeAfter = time.After(time.Duration(r.maxWaitSecond) * time.Second)
	}

	ticker := time.NewTicker(time.Duration(500) * time.Millisecond)
	defer ticker.Stop()

	for range ticker.C {
		select {
		case <-r.ctx.Done():
			return r.ctx.Err()
		case <-timeAfter:
			return errors.New("timeout")
		default:
		}

		err := r.tryLock()

		if err == nil {
			return nil
		}

		if !isErrorNotMyLock(err) {
			return err
		}
	}

	return nil
}

func (r *RedisLockNew) tryLock() error {
	result, err := r.client.SetNX(r.ctx, r.getLockKey(), r.token,
		time.Duration(r.RedisLockOptions.expireSecond)*time.Second).Result()

	if err != nil {
		return err
	}

	if !result {
		return fmt.Errorf("try lock result: %v, err: %w", result, ErrorNotMyLock)
	}

	return nil
}

func (r *RedisLockNew) expireLock() error {
	result, err := r.client.Eval(r.ctx, LuaExpireLock, []string{r.getLockKey()},
		r.token, time.Duration(r.expireSecond)*time.Second).Result()
	if err != nil {
		return err
	}

	if res, ok := result.(int64); ok && res == 1 {
		return nil
	} else {
		return fmt.Errorf("unlock err: ok: %v, res:%v, raw:%v", ok, err, result)
	}
}

func (r *RedisLockNew) runWatchDog(ctx context.Context) {
	ticker := time.NewTicker(time.Duration(10) * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		select {
		case <-ctx.Done():
			return
		default:
		}

		if err := r.expireLock(); err != nil {
			log.Printf("expire lock err: %v", err)
		}
	}
}

func (r *RedisLockNew) watchDog() {
	if !r.isReNew {
		return
	}

	for !atomic.CompareAndSwapInt32(&r.DogRunning, 0, 1) {
	}

	var ctx context.Context
	ctx, r.stopDog = context.WithCancel(r.ctx)
	go func() {
		defer func() {
			atomic.StoreInt32(&r.DogRunning, 0)
		}()

		r.runWatchDog(ctx)
	}()
}

func (r *RedisLockNew) Lock() (err error) {
	defer func() {
		if err == nil {
			r.watchDog()
		}
	}()

	err = r.tryLock()
	if err == nil {
		return nil
	}

	if !isErrorNotMyLock(err) || !r.isBlock {
		return err
	}

	err = r.blockingLock()

	return err
}

func (r *RedisLockNew) UnLock() error {
	defer func() {
		if r.stopDog != nil {
			r.stopDog()
		}
	}()

	result, err := r.client.Eval(r.ctx, LuaUnLock, []string{r.getLockKey()}, r.token).Result()
	if err != nil {
		return err
	}

	if res, ok := result.(int64); ok && res == 1 {
		return nil
	} else {
		return fmt.Errorf("unlock err: ok: %v, res:%v, raw:%v", ok, err, result)
	}
}

func NewRedisLock(client *redis.Client, ctx context.Context, key string, opts ...RedisLockOption) *RedisLockNew {
	lock := &RedisLockNew{
		client: client,
		ctx:    ctx,
		key:    key,
		token:  utils.GetProcessAndGoroutineIDStr(),
	}

	for _, o := range opts {
		o(&lock.RedisLockOptions)
	}

	lock.repairOption()

	return lock
}