go实现分布式锁

go实现分布式锁

package lock


type lock struct {
	key         string
	requestId   string
	checkCancel chan bool
	store       *rds.Client
	expiration  time.Duration
	mu          sync.Mutex
}

// NewLock 初始化分布式锁实例
func NewLock(name, key string, expiration time.Duration) *lock {
	return &lock{store: redis.Client(name), key: key, expiration: expiration, requestId: uuid.New().String(), mu: sync.Mutex{}}
}

// Acquire 获取锁
func (lk *lock) Acquire() bool {
	cxt, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	ok, err := lk.store.SetNX(cxt, lk.key, lk.requestId, lk.expiration).Result()
	if err != nil {
		return false
	}
	if ok {
		// 锁续期检查
		go lk.checkLockIsRelease()
	}
	return ok
}

// checkLockIsRelease 检查锁是否被释放，未被释放就延长锁时间
func (lk *lock) checkLockIsRelease() {
	for {
		checkCxt, _ := context.WithTimeout(context.Background(), time.Millisecond*time.Duration(lk.expiration.Milliseconds()-lk.expiration.Milliseconds()/10))
		lk.checkCancel = make(chan bool)
		select {
		case <-checkCxt.Done():
			// 多次续期，直到锁被释放
			isContinue := lk.done()
			if !isContinue {
				return
			}
		// 取消
		case <-lk.checkCancel:
			return
		}
	}
}

// done 判断锁是否已被释放
func (lk *lock) done() bool {
	lk.mu.Lock()
	defer lk.mu.Unlock()
	cxt, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	res, err := lk.store.Exists(cxt, lk.key).Result()
	cancel()
	if err != nil {
		return false
	}
	if res == 1 {
		cxt, cancel := context.WithTimeout(context.Background(), 3*time.Second)
		ok, err := lk.store.Expire(cxt, lk.key, lk.expiration).Result()
		cancel()
		if err != nil {
			return false
		}
		if ok {
			return true
		}
	}
	return false
}

// Block 阻塞获取锁
func (lk *lock) Block(expiration time.Duration) bool {
	t := time.Now()
	for {
		cxt, cancel := context.WithTimeout(context.Background(), 3*time.Second)
		ok, err := lk.store.SetNX(cxt, lk.key, lk.requestId, lk.expiration).Result()
		cancel()
		if err != nil {
			return false
		}
		if ok {

			go lk.checkLockIsRelease()
			return true
		}
		time.Sleep(200 * time.Millisecond)
		if time.Now().Sub(t) > expiration {
			return false
		}
	}
}

// ForceRelease 强制释放锁，忽略请求id
func (lk *lock) ForceRelease() error {
	lk.mu.Lock()
	defer func() {
		lk.mu.Unlock()
		if lk.checkCancel != nil {
			lk.checkCancel <- true
		}
	}()

	cxt, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	_, err := lk.store.Del(cxt, lk.key).Result()
	return err

}

// Release 释放锁
func (lk *lock) Release() error {
	lk.mu.Lock()
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	rsp, err := lk.store.Eval(ctx, luaScript, []string{lk.key}, lk.requestId).Result()
	lk.mu.Unlock()
	if rsp.(int64) != 0 {
		lk.checkCancel <- true
	}
	return err
}

const luaScript = `
	if redis.call('get', KEYS[1])==ARGV[1] then
		return redis.call('del', KEYS[1])
	else
		return 0
	end
	`

• 使用

l := lock.NewLock(r.RedisClient.GetRedisName(), r.RedisClient.GenKey4LockRecommend(now.Format("2006-01-02")), 30*time.Second)
if !l.Acquire() {
		logs.Info(ctx, "createRecommendInstanceAfterAb", zap.Any("acquire", ""))
		return
	}
	defer func() {
		er := l.Release()
		if er != nil {
			logs.Error(ctx, "redis lock release failed", zap.Error(er))
		}
	}()