base-golang/pkg/limiter/ip.go

package limiter

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

var _ IPRateLimiter = (*ipRateLimiter)(nil)

type IPRateLimiter interface {
	addIP(ip string) *rate.Limiter
	GetLimiter(ip string) *rate.Limiter
}

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

func NewIPRateLimiter(limit time.Duration, burst int) IPRateLimiter {
	return &ipRateLimiter{
		ips:   sync.Map{},
		limit: rate.Every(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)
}