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 := &item{ lastTime: time.Now(), limiter: rate.NewLimiter(i.limit, i.burst), } i.list.Store(key, store) return store.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() }