first commit
This commit is contained in:
71
pkg/limiter/rate.go
Normal file
71
pkg/limiter/rate.go
Normal file
@ -0,0 +1,71 @@
|
||||
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()
|
||||
}
|
Reference in New Issue
Block a user