gd_auth_check/limit/rate.go

package limit

import (
	"gd_auth_check/common.in/config"
	"sync"
	"time"
)

const (
	releaseInterval = time.Second
)

type Limiter struct {
	mu sync.Mutex
	// 总容量
	burst int
	// 已使用数
	used int
	// 已使用的令牌
	usedToken map[string]int64
	// 时间轮，用以自动清理令牌
	tiker *time.Ticker
}

func NewLimiter(burst int) *Limiter {
	return &Limiter{
		burst:     burst,
		usedToken: make(map[string]int64),
	}
}

func (lim *Limiter) Allow() (bool, string) {
	ok, tokenArr := lim.reserveN(1)
	if ok {
		return ok, tokenArr[0]
	}
	return ok, ""
}

func (lim *Limiter) AllowN(n int) (bool, []string) {
	return lim.reserveN(n)
}

func (lim *Limiter) reserveN(n int) (bool, []string) {
	lim.mu.Lock()
	defer lim.mu.Unlock()

	tokenArr := make([]string, 0, n)
	if lim.burst-lim.used-n >= 0 {
		lim.used += n

		for i := 0; i < n; i++ {
			// 生成令牌
			token := getToken(fromMem)
			tokenArr = append(tokenArr, token)

			// 存储令牌
			lim.usedToken[token] = time.Now().Unix()
		}

		go lim.ticker()

		return true, tokenArr
	}

	return false, nil
}

func (lim *Limiter) releaseToken(token string) {
	lim.mu.Lock()
	defer lim.mu.Unlock()

	delete(lim.usedToken, token)
	lim.used--
}

func (lim *Limiter) ticker() {
	if lim.tiker == nil {
		lim.tiker = time.NewTicker(releaseInterval)
	}

	defer func() {
		lim.tiker.Stop()
	}()

	if len(lim.usedToken) == 0 {
		return
	}

	interval, _ := config.Conf.RateLimit.RescueTickerTime.Int64()
	for range lim.tiker.C {
		if len(lim.usedToken) == 0 {
			break
		}

		ts := time.Now().Unix()
		for k, v := range lim.usedToken {
			// 超过3秒的则直接释放
			if ts-v > interval {
				delete(lim.usedToken, k)
				lim.used--
			}
		}

		// 全部释放完
		if lim.used == 0 {
			break
		}
	}
}