s_cherry_e/sr/internal/component/rate/func_limiter.go

package rate

import (
	"f1-game/internal/code"
)

type (
	FuncLimiter struct {
		entry                         // 全局的默认设置
		funcEntries map[string]*entry // key: funcName,value: entry
	}

	entry struct {
		interval   int64 // 访问时间间隔(ms)
		visitLimit int32 // interval时间内最大访问次数
		visitNum   int32 // 已访问次数
		errCode    int32 // 返回错误码
		deadline   int64 // 截止时间(ms)
	}
)

func NewFuncLimiter(defaultInterval int64, defaultVisitLimit int32, defaultCode int32) FuncLimiter {
	return FuncLimiter{
		entry:       newEntry(defaultInterval, defaultVisitLimit, defaultCode),
		funcEntries: map[string]*entry{},
	}
}

func newEntry(interval int64, visitLimit int32, errCode int32) entry {
	return entry{
		interval:   interval,
		visitLimit: visitLimit,
		visitNum:   0,
		errCode:    errCode,
		deadline:   0,
	}
}

func (p *FuncLimiter) Add(funcName string, interval int64, visitLimit int32, errCode int32) {
	entry := newEntry(interval, visitLimit, errCode)
	p.funcEntries[funcName] = &entry
}

func (p *FuncLimiter) Check(now int64, funcName string) int32 {
	if entry, found := p.funcEntries[funcName]; found {
		return entry.check(now)
	}

	return p.check(now)
}

func (p *entry) check(now int64) int32 {
	if p.deadline >= now {
		p.visitNum++

		if p.visitNum >= p.visitLimit {
			return p.errCode
		}

		return code.OK
	}

	p.deadline = now + p.interval
	p.visitNum = 0

	return code.OK
}