ginuerzh
9bbd7da526
- limiterGroup.Allow(): return false immediately on first denial instead of continuing through all limiters (wasting tokens on subsequent limiters). Sort limiters by Limit() so most restrictive fires first. - rateLimiter.Close(): add missing httpLoader.Close() to prevent resource leak. - Add doc comments for all exported symbols (package, constants, Options, NewRateLimiter, NewLimiter, ErrRateLimit, generators). - Add rate_test.go with 59 unit tests covering generators, limiter groups, parsing (limit/line/patterns), reload, loader merge, Close behavior, and priority ordering (exact IP > CIDR > $$ > $).
62 lines
1.2 KiB
Go
62 lines
1.2 KiB
Go
package rate
|
|
|
|
import (
|
|
"errors"
|
|
"sort"
|
|
"time"
|
|
|
|
limiter "github.com/go-gost/core/limiter/rate"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// ErrRateLimit is returned when a rate limit is exceeded.
|
|
var ErrRateLimit = errors.New("rate limit")
|
|
|
|
type rlimiter struct {
|
|
limiter *rate.Limiter
|
|
}
|
|
|
|
// NewLimiter creates a [limiter.Limiter] with the specified rate r (events per second)
|
|
// and burst b.
|
|
func NewLimiter(r float64, b int) limiter.Limiter {
|
|
return &rlimiter{
|
|
limiter: rate.NewLimiter(rate.Limit(r), b),
|
|
}
|
|
}
|
|
|
|
func (l *rlimiter) Allow(n int) bool {
|
|
return l.limiter.AllowN(time.Now(), n)
|
|
}
|
|
|
|
func (l *rlimiter) Limit() float64 {
|
|
return float64(l.limiter.Limit())
|
|
}
|
|
|
|
type limiterGroup struct {
|
|
limiters []limiter.Limiter
|
|
}
|
|
|
|
func newLimiterGroup(limiters ...limiter.Limiter) *limiterGroup {
|
|
sort.Slice(limiters, func(i, j int) bool {
|
|
return limiters[i].Limit() < limiters[j].Limit()
|
|
})
|
|
return &limiterGroup{limiters: limiters}
|
|
}
|
|
|
|
func (l *limiterGroup) Allow(n int) bool {
|
|
for i := range l.limiters {
|
|
if !l.limiters[i].Allow(n) {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func (l *limiterGroup) Limit() float64 {
|
|
if len(l.limiters) == 0 {
|
|
return 0
|
|
}
|
|
|
|
return l.limiters[0].Limit()
|
|
}
|