base-golang/pkg/limiter/rate.go

package limiter

import (
	"git.bvbej.com/bvbej/base-golang/pkg/ticker"
	"golang.org/x/time/rate"
	"sync"
	"time"
)

var _ RateLimiter = (*rateLimiter)(nil)

type item struct {
	lastTime time.Time
	limiter  *rate.Limiter
}

type RateLimiter interface {
	set(key string) *rate.Limiter
	get(key string) *rate.Limiter
	Allow(key string) bool
}

type rateLimiter struct {
	limit   rate.Limit
	burst   int
	list    *sync.Map
	recycle ticker.Ticker
}

func NewRateLimiter(limit rate.Limit, burst int) RateLimiter {
	list := new(sync.Map)
	t := ticker.New(time.Minute)
	t.Process(func() {
		list.Range(func(key, value any) bool {
			if value.(*item).lastTime.Before(time.Now().Add(-time.Hour)) {
				list.Delete(key)
			}
			return true
		})
	})
	return &rateLimiter{
		list:    list,
		limit:   limit,
		recycle: t,
		burst:   burst,
	}
}

func (i *rateLimiter) set(key string) *rate.Limiter {
	store, _ := i.list.LoadOrStore(key, &item{
		lastTime: time.Now(),
		limiter:  rate.NewLimiter(i.limit, i.burst),
	})
	return store.(*item).limiter
}

func (i *rateLimiter) get(key string) *rate.Limiter {
	value, ok := i.list.Load(key)
	if !ok {
		return i.set(key)
	}
	value.(*item).lastTime = time.Now()
	i.list.Store(key, value)
	return value.(*item).limiter
}

func (i *rateLimiter) Allow(key string) bool {
	limiter := i.get(key)
	return limiter.Allow()
}