base-golang/pkg/limiter/rate.go

package limiter

import (
	"golang.org/x/time/rate"
	"sync"
)

var _ RateLimiter = (*rateLimiter)(nil)

type RateLimiter interface {
	addKey(key string) *rate.Limiter
	getLimiter(key string) *rate.Limiter
	Allow(key string) bool
}

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

func NewRateLimiter(limit int, burst int) RateLimiter {
	return &rateLimiter{
		list:  new(sync.Map),
		limit: rate.Limit(limit),
		burst: burst,
	}
}

func (i *rateLimiter) addKey(key string) *rate.Limiter {
	store, _ := i.list.LoadOrStore(key, rate.NewLimiter(i.limit, i.burst))
	return store.(*rate.Limiter)
}

func (i *rateLimiter) getLimiter(key string) *rate.Limiter {
	limiter, exists := i.list.Load(key)
	if !exists {
		return i.addKey(key)
	}
	return limiter.(*rate.Limiter)
}

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