base-golang/pkg/limiter/rate.go

72 lines
1.3 KiB
Go
Raw Normal View History

2024-07-23 10:23:43 +08:00
package limiter
import (
2024-07-31 16:49:14 +08:00
"gitea.bvbej.com/bvbej/base-golang/pkg/ticker"
2024-07-23 10:23:43 +08:00
"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()
}