openbilibili/library/net/http/blademaster/middleware/rate/limit.go

package rate

import (
	"net/http"
	"sync/atomic"

	"go-common/library/log"
	bm "go-common/library/net/http/blademaster"

	"golang.org/x/time/rate"
)

const (
	_defBurst = 100
)

// Limiter controls how frequently events are allowed to happen.
type Limiter struct {
	apps atomic.Value
	urls atomic.Value
}

// Config limitter  conf.
type Config struct {
	Apps map[string]*Limit
	URLs map[string]*Limit
}

// Limit limit conf.
type Limit struct {
	Limit rate.Limit
	Burst int
}

// New return Limiter.
func New(conf *Config) (l *Limiter) {
	l = &Limiter{}
	l.apps.Store(make(map[string]*rate.Limiter))
	l.urls.Store(make(map[string]*rate.Limiter))
	if conf != nil {
		l.Reload(conf)
	}
	return
}

// Reload reload limit conf.
func (l *Limiter) Reload(c *Config) {
	if c == nil {
		return
	}
	var (
		ok  bool
		al  *rate.Limiter
		ul  *rate.Limiter
		as  map[string]*rate.Limiter
		nas map[string]*rate.Limiter
		us  map[string]*rate.Limiter
		nus map[string]*rate.Limiter
	)
	if as, ok = l.apps.Load().(map[string]*rate.Limiter); !ok {
		log.Error("apps limiter load map hava no data ")
		return
	}
	nas = make(map[string]*rate.Limiter, len(as))
	for k, v := range as {
		nas[k] = v
	}
	for k, v := range c.Apps {
		if al, ok = nas[k]; !ok || (al.Burst() != v.Burst || al.Limit() != v.Limit) {
			nas[k] = rate.NewLimiter(v.fix())
		}
	}
	l.apps.Store(nas)

	if us, ok = l.urls.Load().(map[string]*rate.Limiter); !ok {
		log.Error("urls limiter load map hava no data ")
		return
	}
	nus = make(map[string]*rate.Limiter, len(us))
	for k, v := range us {
		nus[k] = v
	}
	for k, v := range c.URLs {
		if ul, ok = nus[k]; !ok || (ul.Burst() != v.Burst || ul.Limit() != v.Limit) {
			nus[k] = rate.NewLimiter(v.fix())
		}
	}
	l.urls.Store(nus)
}

func (l *Limit) fix() (lim rate.Limit, b int) {
	lim = rate.Inf
	b = _defBurst
	if l.Limit <= 0 {
		lim = rate.Inf
	} else {
		lim = l.Limit
	}
	if l.Burst > 0 {
		b = l.Burst
	}
	return
}

// Allow reports whether event may happen at time now.
func (l *Limiter) Allow(appKey, path string) bool {
	if as, ok := l.apps.Load().(map[string]*rate.Limiter); ok {
		if lim, ok := as[appKey]; ok {
			if !lim.Allow() {
				return false
			}
		}
	}
	if us, ok := l.urls.Load().(map[string]*rate.Limiter); ok {
		if lim, ok := us[path]; ok {
			if !lim.Allow() {
				return false
			}
		}
	}
	return true
}

func (l *Limiter) ServeHTTP(c *bm.Context) {
	req := c.Request
	appkey := req.Form.Get("appkey")
	path := req.URL.Path
	if !l.Allow(appkey, path) {
		c.AbortWithStatus(http.StatusTooManyRequests)
		return
	}
}

// Handler is router allow handle.
func (l *Limiter) Handler() bm.HandlerFunc {
	return l.ServeHTTP
}