gd_auth_check/common.in/utils/distribute_lock.go

package utils

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

	"gd_auth_check/common.in/cache"
	"go.etcd.io/etcd/client"
)

var lockKey = "/gd/lock"
var lockValue = "locked"

var etcdClient client.Client

var keyApi client.KeysAPI
var gmutex sync.Mutex
var gm = map[string]string{}

func SetEtcdKeyApi(c client.Client) {
	if keyApi == nil {
		keyApi = client.NewKeysAPI(c)
	}
}

func setKey(keyApi client.KeysAPI, key string, value string, prevExist client.PrevExistType) error {
	opts := &client.SetOptions{
		PrevExist: prevExist,
		TTL:       10 * time.Second,
	}
	_, err := keyApi.Set(context.Background(), key, value, opts)
	return err
}

func delKey(keyApi client.KeysAPI, key string) error {
	_, err := keyApi.Delete(context.Background(), key, nil)
	e, _ := err.(client.Error)
	if e.Code == client.ErrorCodeKeyNotFound {
		return nil
	}
	return err
}

func lock(realLockKey string) error {
	var err error
	err = setKey(keyApi, realLockKey, lockValue, client.PrevNoExist)
	if err == nil {
		return nil
	}
	return err
}

func redisLock(key string) error {
	v, err := cache.Redis.SetNxEx(key, "-", 10)
	if v == false {
		return errors.New("locked:" + key)
	}
	return err
}

func redisUnlock(key string) error {
	_, err := cache.Redis.Del(key)
	return err
}

func localLock(key string) bool {
	now := time.Now().Unix()
	for {
		gmutex.Lock()
		if _, ok := gm[key]; ok == false {
			gm[key] = "-"
			gmutex.Unlock()
			return true
		}
		gmutex.Unlock()
		if time.Now().Unix()-now >= 60 {
			return false
		}
		time.Sleep(50 * time.Millisecond)
	}
}

func localUnlock(key string) {
	gmutex.Lock()
	if _, ok := gm[key]; ok == true {
		delete(gm, key)
	}
	if len(gm) == 0 {
		gm = map[string]string{}
	}
	gmutex.Unlock()
}

func Lock(flag string) error {
	realLockKey := lockKey + "/" + flag
	if localLock(realLockKey) == false {
		return errors.New("local lock timeout")
	}
	count := 0
	for {
		count++
		if count > 1000 {
			localUnlock(realLockKey)
			return errors.New("max count")
		}
		err := redisLock(realLockKey)
		if err == nil {
			return nil
		}
		time.Sleep(400 * time.Millisecond)
	}
	return nil
}

func UnLock(flag string) error {
	err := redisUnlock(lockKey + "/" + flag)
	localUnlock(lockKey + "/" + flag)
	return err
}