woop/internal/rate/limiter.go

package rate

import (
	"container/list"
	"context"
	"sync"
	"time"

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

type (
	limiterPerID struct {
		evictList *list.List
		items     map[string]*list.Element

		lock sync.Mutex
		t    timer
		c    configInt
	}

	entry struct {
		id       string
		limiter  *rate.Limiter
		lastCall time.Time
	}
)

// NewLimiterPerID creates a limit limiter based on ID
func NewLimiterPerID(rate rate.Limit, burst int, c *Config) *limiterPerID {
	lpi := &limiterPerID{
		evictList: list.New(),
		items:     make(map[string]*list.Element),
		t:         timerImpl{},
		c:         toConfigInt(rate, burst, c),
	}
	go lpi.maintainLoop()
	return lpi
}

// AllowN return whether a request of size n is allowed by IP
func (lpi *limiterPerID) AllowN(id string, n int) bool {
	if lpi.isWhitelisted(id) {
		return true
	}
	lmt := lpi.getLimiterByID(id)
	return lmt.AllowN(lpi.t.now(), n)
}

// WaitN waits until the n token is granted by IP
func (lpi *limiterPerID) WaitN(ctx context.Context, id string, n int) error {
	if lpi.isWhitelisted(id) {
		return nil
	}
	lmt := lpi.getLimiterByID(id)
	return lmt.WaitN(ctx, n)
}

func (lpi *limiterPerID) getLimiterByID(id string) *rate.Limiter {
	lpi.lock.Lock()
	defer lpi.lock.Unlock()

	elem, ok := lpi.items[id]
	if !ok || elem == nil {
		lmt := rate.NewLimiter(lpi.c.limit, lpi.c.burst)
		ent := &entry{id: id, limiter: lmt, lastCall: lpi.t.now()}
		elem := lpi.evictList.PushFront(ent)
		lpi.items[id] = elem
		return lmt
	}
	ent := elem.Value.(*entry)
	lpi.evictList.MoveToFront(elem)
	ent.lastCall = lpi.t.now()
	return ent.limiter
}

func (lpi *limiterPerID) maintainLoop() {
	t := lpi.t.newTicker(lpi.c.checkInt)
	defer t.Stop()

	for range t.C {
		lpi.maintain()
	}
}

func (lpi *limiterPerID) maintain() {
	lpi.lock.Lock()
	defer lpi.lock.Unlock()

	for lpi.evictList.Len() > lpi.c.capacity {
		evicted := lpi.evictLast(func(ent *entry) bool {
			return lpi.t.since(ent.lastCall) > lpi.c.minEvictDur
		})
		if !evicted {
			break
		}
	}
}

func (lpi *limiterPerID) evictLast(isEvict func(ent *entry) bool) bool {
	elem := lpi.evictList.Back()
	evict := elem.Value.(*entry) // not nil ensured
	if isEvict(evict) {
		delete(lpi.items, evict.id)
		lpi.evictList.Remove(elem)
		return true
	}
	return false
}

func (lpi *limiterPerID) isWhitelisted(id string) bool {
	_, ok := lpi.c.whitelist[id]
	return ok
}