ginuerzh
161 lines
3.1 KiB
Go
161 lines
3.1 KiB
Go
package limiter
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/go-gost/core/limiter"
|
|
"github.com/go-gost/core/limiter/traffic"
|
|
)
|
|
|
|
const (
|
|
defaultRefreshInterval = 30 * time.Second
|
|
defaultCleanupInterval = 60 * time.Second
|
|
)
|
|
|
|
type options struct {
|
|
refreshInterval time.Duration
|
|
cleanupInterval time.Duration
|
|
scope string
|
|
}
|
|
|
|
type Option func(*options)
|
|
|
|
func RefreshIntervalOption(interval time.Duration) Option {
|
|
return func(o *options) {
|
|
o.refreshInterval = interval
|
|
}
|
|
}
|
|
|
|
func CleanupIntervalOption(interval time.Duration) Option {
|
|
return func(o *options) {
|
|
o.cleanupInterval = interval
|
|
}
|
|
}
|
|
|
|
func ScopeOption(scope string) Option {
|
|
return func(o *options) {
|
|
o.scope = scope
|
|
}
|
|
}
|
|
|
|
type cachedTrafficLimiter struct {
|
|
inLimits *Cache
|
|
outLimits *Cache
|
|
limiter traffic.TrafficLimiter
|
|
options options
|
|
}
|
|
|
|
func NewCachedTrafficLimiter(limiter traffic.TrafficLimiter, opts ...Option) traffic.TrafficLimiter {
|
|
if limiter == nil {
|
|
return nil
|
|
}
|
|
|
|
var options options
|
|
for _, opt := range opts {
|
|
opt(&options)
|
|
}
|
|
|
|
if options.refreshInterval == 0 {
|
|
options.refreshInterval = defaultRefreshInterval
|
|
}
|
|
if options.refreshInterval < time.Second {
|
|
options.refreshInterval = time.Second
|
|
}
|
|
|
|
if options.cleanupInterval == 0 {
|
|
options.cleanupInterval = defaultCleanupInterval
|
|
}
|
|
if options.cleanupInterval < time.Second {
|
|
options.cleanupInterval = time.Second
|
|
}
|
|
|
|
lim := &cachedTrafficLimiter{
|
|
inLimits: NewCache(options.cleanupInterval),
|
|
outLimits: NewCache(options.cleanupInterval),
|
|
limiter: limiter,
|
|
options: options,
|
|
}
|
|
return lim
|
|
}
|
|
|
|
func (p *cachedTrafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
|
|
if p.limiter == nil {
|
|
return nil
|
|
}
|
|
|
|
var options limiter.Options
|
|
for _, opt := range opts {
|
|
opt(&options)
|
|
}
|
|
|
|
if p.options.scope != "" && p.options.scope != options.Scope {
|
|
return nil
|
|
}
|
|
|
|
item := p.inLimits.Get(key)
|
|
lim, _ := item.Value().(traffic.Limiter)
|
|
if !item.Expired() {
|
|
return lim
|
|
}
|
|
|
|
limNew := p.limiter.In(ctx, key, opts...)
|
|
if limNew == nil {
|
|
limNew = lim
|
|
}
|
|
if item == nil || !p.equal(lim, limNew) {
|
|
p.inLimits.Set(key, NewItem(limNew, p.options.refreshInterval))
|
|
return limNew
|
|
}
|
|
|
|
p.inLimits.Set(key, NewItem(lim, p.options.refreshInterval))
|
|
|
|
return lim
|
|
}
|
|
|
|
func (p *cachedTrafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
|
|
if p.limiter == nil {
|
|
return nil
|
|
}
|
|
|
|
var options limiter.Options
|
|
for _, opt := range opts {
|
|
opt(&options)
|
|
}
|
|
|
|
if p.options.scope != "" && p.options.scope != options.Scope {
|
|
return nil
|
|
}
|
|
|
|
item := p.outLimits.Get(key)
|
|
lim, _ := item.Value().(traffic.Limiter)
|
|
if !item.Expired() {
|
|
return lim
|
|
}
|
|
|
|
limNew := p.limiter.Out(ctx, key, opts...)
|
|
if limNew == nil {
|
|
limNew = lim
|
|
}
|
|
if item == nil || !p.equal(lim, limNew) {
|
|
p.outLimits.Set(key, NewItem(limNew, p.options.refreshInterval))
|
|
return limNew
|
|
}
|
|
|
|
p.outLimits.Set(key, NewItem(lim, p.options.refreshInterval))
|
|
|
|
return lim
|
|
}
|
|
|
|
func (p *cachedTrafficLimiter) equal(lim1, lim2 traffic.Limiter) bool {
|
|
if lim1 == lim2 {
|
|
return true
|
|
}
|
|
|
|
if lim1 == nil || lim2 == nil {
|
|
return false
|
|
}
|
|
|
|
return lim1.Limit() == lim2.Limit()
|
|
}
|