go-mixed/distributed_locker/distributed_lock.go

package locker

import (
	"context"
	"errors"
	"strings"
	"sync"
	"time"

	"github.com/charlienet/go-mixed/rand"
	"github.com/charlienet/go-mixed/redis"
	goredis "github.com/redis/go-redis/v9"
)

const (
	// 加锁(可重入)
	lockCmd = `if redis.call("GET", KEYS[1]) == ARGV[1] then
	redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2])
	return "OK"
else
	return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2])
end`

	// 解锁
	delCmd = `if redis.call("GET", KEYS[1]) == ARGV[1] then
	return redis.call("DEL", KEYS[1])
else
	return '0'
end`

	// 延期
	incrCmd = `
if redis.call('get', KEYS[1]) == ARGV[1] then
  return redis.call('expire', KEYS[1], ARGV[2])
 else
   return '0'
end`
)

const (
	defaultExpire = time.Second * 10
	retryInterval = time.Millisecond * 10
)

var (
	once             sync.Once
	ErrContextCancel = errors.New("context cancel")
)

type distributedlock struct {
	clients  []redis.Client  // redis 客户端
	ctx      context.Context //
	key      string          // 资源键
	rand     string          // 随机值
	unlocked bool            // 是否已解锁
	expire   time.Duration   // 过期时间
}

func NewDistributedLocker(ctx context.Context, key string, clients ...redis.Client) *distributedlock {
	expire := defaultExpire
	if deadline, ok := ctx.Deadline(); ok {
		expire = time.Until(deadline)
	}

	locker := &distributedlock{
		ctx:     ctx,
		clients: clients,
		key:     key,
		rand:    rand.Hex.Generate(24),
		expire:  expire,
	}

	return locker
}

func (locker *distributedlock) Lock() error {
	for {
		select {
		case <-locker.ctx.Done():
			return ErrContextCancel
		default:
			if locker.TryLock() {
				return nil
			}
		}

		time.Sleep(retryInterval)
	}
}

func (locker *distributedlock) TryLock() bool {
	results := locker.Eval(locker.ctx, lockCmd, []string{locker.key}, locker.rand, locker.expire.Milliseconds())

	if !isSuccess(results) {
		locker.Unlock()
		return false
	}

	locker.expandLockTime()

	return true
}

func (locker *distributedlock) Unlock() {
	locker.Eval(locker.ctx, delCmd, []string{locker.key}, locker.rand)
	locker.unlocked = true
}

func (l *distributedlock) expandLockTime() {
	once.Do(func() {
		go func() {
			for {
				time.Sleep(l.expire / 3)
				if l.unlocked {
					break
				}

				l.resetExpire()
			}
		}()
	})
}

func (locker *distributedlock) resetExpire() {
	locker.Eval(locker.ctx,
		incrCmd,
		[]string{locker.key},
		locker.rand,
		locker.expire.Seconds())
}

func (locker *distributedlock) Eval(ctx context.Context, cmd string, keys []string, args ...any) []*goredis.Cmd {
	results := make([]*goredis.Cmd, 0, len(locker.clients))

	var wg sync.WaitGroup
	wg.Add(len(locker.clients))
	for _, rdb := range locker.clients {
		go func(rdb redis.Client) {
			defer wg.Done()

			results = append(results, rdb.Eval(ctx, cmd, keys, args...))
		}(rdb)
	}

	wg.Wait()

	return results
}

func isSuccess(results []*goredis.Cmd) bool {
	successCount := 0
	for _, ret := range results {
		resp, err := ret.Result()

		if err != nil || resp == nil {
			return false
		}

		reply, ok := resp.(string)
		if ok && strings.EqualFold(reply, "OK") {
			successCount++
		}
	}

	return successCount >= len(results)/2+1
}