go-common/lockx/lock.go

package lockx

import (
	"context"
	"github.com/zeromicro/go-zero/core/stores/redis"
	"time"
)

type DistributedLock struct {
	client  *redis.Redis
	key     string
	seconds int
}

func NewDistributedLock(client *redis.Redis, key string, seconds int) *DistributedLock {
	return &DistributedLock{
		client:  client,
		key:     key,
		seconds: seconds,
	}
}

// Lock 尝试获取分布式锁
func (dl *DistributedLock) Lock() (bool, error) {
	if ok, err := dl.client.SetnxEx(dl.key, "1", dl.seconds); err != nil {
		return false, err
	} else if ok {
		return true, nil
	}
	return false, nil
}

// WaitLock 尝试获取分布式锁,获取不到等待
func (dl *DistributedLock) WaitLock(ctx context.Context) error {
	for {
		select {
		case <-ctx.Done():
			// 如果Context被取消，则直接返回错误
			return ctx.Err()
		default:
			// 尝试获取锁
			result, err := dl.client.SetnxEx(dl.key, "1", dl.seconds)
			if err != nil {
				return err
			}

			if result {
				// 获取到锁，直接返回
				return nil
			}

			// 等待一段时间后重试
			time.Sleep(time.Millisecond * 50)
		}
	}
}

// Unlock 释放分布式锁
func (dl *DistributedLock) Unlock() error {
	// 使用 Lua 脚本来保证原子性地删除锁（避免误删其他客户端的锁）
	luaScript := `
		if redis.call("get", KEYS[1]) == ARGV[1] then
			return redis.call("del", KEYS[1])
		else
			return 0
		end
	`
	_, err := dl.client.Eval(luaScript, []string{dl.key}, "1")
	return err
}