45 lines
781 B
Go
45 lines
781 B
Go
package rate
|
|
|
|
import (
|
|
"github.com/go-gost/core/limiter/rate"
|
|
limiter "github.com/go-gost/core/limiter/rate"
|
|
)
|
|
|
|
type RateLimitGenerator interface {
|
|
Limiter() limiter.Limiter
|
|
}
|
|
|
|
type rateLimitGenerator struct {
|
|
r float64
|
|
}
|
|
|
|
func NewRateLimitGenerator(r float64) RateLimitGenerator {
|
|
return &rateLimitGenerator{
|
|
r: r,
|
|
}
|
|
}
|
|
|
|
func (p *rateLimitGenerator) Limiter() limiter.Limiter {
|
|
if p == nil || p.r <= 0 {
|
|
return nil
|
|
}
|
|
return NewLimiter(p.r, int(p.r)+1)
|
|
}
|
|
|
|
type rateLimitSingleGenerator struct {
|
|
limiter rate.Limiter
|
|
}
|
|
|
|
func NewRateLimitSingleGenerator(r float64) RateLimitGenerator {
|
|
p := &rateLimitSingleGenerator{}
|
|
if r > 0 {
|
|
p.limiter = NewLimiter(r, int(r)+1)
|
|
}
|
|
|
|
return p
|
|
}
|
|
|
|
func (p *rateLimitSingleGenerator) Limiter() limiter.Limiter {
|
|
return p.limiter
|
|
}
|