base-golang/pkg/limiter/ip.go

package limiter

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

var _ IPRateLimiter = (*ipRateLimiter)(nil)

type IPRateLimiter interface {
	addIP(ip string) *rate.Limiter
	getLimiter(ip string) *rate.Limiter
	Allow(ip string) bool
}

type ipRateLimiter struct {
	ips   *sync.Map
	limit rate.Limit
	burst int
}

func NewIPRateLimiter(limit int, burst int) IPRateLimiter {
	return &ipRateLimiter{
		ips:   new(sync.Map),
		limit: rate.Limit(limit),
		burst: burst,
	}
}

func (i *ipRateLimiter) addIP(ip string) *rate.Limiter {
	store, _ := i.ips.LoadOrStore(ip, rate.NewLimiter(i.limit, i.burst))
	return store.(*rate.Limiter)
}

func (i *ipRateLimiter) getLimiter(ip string) *rate.Limiter {
	limiter, exists := i.ips.Load(ip)
	if !exists {
		return i.addIP(ip)
	}
	return limiter.(*rate.Limiter)
}

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