add scope param for traffic limiter

internal/util/limiter/cache.go

@@ -0,0 +1,64 @@
+package limiter
+
+import (
+	"sync"
+	"time"
+)
+
+type Item struct {
+	v          interface{}
+	expiration int64
+}
+
+func NewItem(v interface{}, d time.Duration) *Item {
+	var expiration int64
+	if d > 0 {
+		expiration = time.Now().Add(d).UnixNano()
+	}
+
+	return &Item{
+		v:          v,
+		expiration: expiration,
+	}
+}
+
+func (p *Item) Expired() bool {
+	if p == nil {
+		return true
+	}
+	return p.expiration > 0 && time.Now().UnixNano() > p.expiration
+}
+
+func (p *Item) Value() interface{} {
+	if p == nil {
+		return nil
+	}
+	return p.v
+}
+
+type Cache struct {
+	items           map[string]*Item
+	cleanupInterval time.Duration
+	mu              sync.RWMutex
+}
+
+func NewCache(cleanupInterval time.Duration) *Cache {
+	return &Cache{
+		cleanupInterval: cleanupInterval,
+		items:           make(map[string]*Item),
+	}
+}
+
+func (c *Cache) Set(key string, item *Item) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.items[key] = item
+}
+
+func (c *Cache) Get(key string) *Item {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	return c.items[key]
+}

internal/util/limiter/limiter.go

@@ -0,0 +1,92 @@
+package limiter
+
+import (
+	"context"
+	"time"
+
+	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/traffic"
+)
+
+type cachedTrafficLimiter struct {
+	inLimits  *Cache
+	outLimits *Cache
+	limiter   traffic.TrafficLimiter
+	ttl       time.Duration
+}
+
+func NewCachedTrafficLimiter(limiter traffic.TrafficLimiter, ttl time.Duration, cleanupInterval time.Duration) traffic.TrafficLimiter {
+	if limiter == nil {
+		return nil
+	}
+
+	lim := &cachedTrafficLimiter{
+		inLimits:  NewCache(cleanupInterval),
+		outLimits: NewCache(cleanupInterval),
+		limiter:   limiter,
+		ttl:       ttl,
+	}
+	return lim
+}
+
+func (p *cachedTrafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	if p.limiter == nil {
+		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.ttl))
+		return limNew
+	}
+
+	p.inLimits.Set(key, NewItem(lim, p.ttl))
+
+	return lim
+}
+
+func (p *cachedTrafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	if p.limiter == nil {
+		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.ttl))
+		return limNew
+	}
+
+	p.outLimits.Set(key, NewItem(lim, p.ttl))
+
+	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()
+}