ginuerzh
64f9a048f8
- Add NewLimiterWithBurst(rate, burst) for independent burst control - Extend parseLimit to accept optional 4th field as burst size - Add DroppedPacketCounter interface with DroppedPackets() for packet/udp conns - Export ErrRateLimited sentinel for rate-limit error checks - Fix IP-level cache expiration (NoExpiration → defaultExpiration) to prevent unbounded growth; reset TTL on access in In/Out - Short-circuit single-limiter path to avoid limiterGroup allocation - Add nil guard in cachedTrafficLimiter when both old and new values are nil - Remove unused ctx field from wrapper structs
92 lines
1.8 KiB
Go
92 lines
1.8 KiB
Go
package traffic
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sort"
|
|
"strconv"
|
|
|
|
limiter "github.com/go-gost/core/limiter/traffic"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
type llimiter struct {
|
|
limiter *rate.Limiter
|
|
}
|
|
|
|
// NewLimiter creates a rate-based traffic Limiter with the given rate
|
|
// (bytes per second). The burst size is set equal to the rate.
|
|
func NewLimiter(r int) limiter.Limiter {
|
|
return &llimiter{
|
|
limiter: rate.NewLimiter(rate.Limit(r), r),
|
|
}
|
|
}
|
|
|
|
// NewLimiterWithBurst creates a rate-based traffic Limiter with the given rate
|
|
// (bytes per second) and burst size. If burst <= 0, it defaults to r.
|
|
func NewLimiterWithBurst(r, burst int) limiter.Limiter {
|
|
if burst <= 0 {
|
|
burst = r
|
|
}
|
|
return &llimiter{
|
|
limiter: rate.NewLimiter(rate.Limit(r), burst),
|
|
}
|
|
}
|
|
|
|
func (l *llimiter) Wait(ctx context.Context, n int) int {
|
|
if l.limiter.Burst() < n {
|
|
n = l.limiter.Burst()
|
|
}
|
|
if err := l.limiter.WaitN(ctx, n); err != nil {
|
|
return 0
|
|
}
|
|
return n
|
|
}
|
|
|
|
func (l *llimiter) Limit() int {
|
|
return int(l.limiter.Limit())
|
|
}
|
|
|
|
func (l *llimiter) Set(n int) {
|
|
l.limiter.SetLimit(rate.Limit(n))
|
|
l.limiter.SetBurst(n)
|
|
}
|
|
|
|
func (l *llimiter) String() string {
|
|
return strconv.Itoa(int(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) Wait(ctx context.Context, n int) int {
|
|
for i := range l.limiters {
|
|
if v := l.limiters[i].Wait(ctx, n); v < n {
|
|
n = v
|
|
}
|
|
}
|
|
return n
|
|
}
|
|
|
|
func (l *limiterGroup) Limit() int {
|
|
if len(l.limiters) == 0 {
|
|
return 0
|
|
}
|
|
|
|
return l.limiters[0].Limit()
|
|
}
|
|
|
|
func (l *limiterGroup) Set(n int) {}
|
|
|
|
func (l *limiterGroup) String() string {
|
|
return fmt.Sprintf("%v", l.limiters)
|
|
}
|