From 15feb7599e325772fecd5745b0d5581bd89e094e Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Mon, 19 Dec 2022 19:33:29 +0800 Subject: [PATCH] fix traffic limiter --- go.mod | 1 + go.sum | 2 + limiter/traffic/generator.go | 40 +-- limiter/traffic/limiter.go | 48 +++- limiter/traffic/traffic.go | 433 ++++++++++++++++++++++---------- limiter/traffic/wrapper/conn.go | 160 ++++++------ 6 files changed, 438 insertions(+), 246 deletions(-) diff --git a/go.mod b/go.mod index 8683f24..6026b29 100644 --- a/go.mod +++ b/go.mod @@ -18,6 +18,7 @@ require ( github.com/gorilla/websocket v1.5.0 github.com/lucas-clemente/quic-go v0.30.0 github.com/miekg/dns v1.1.50 + github.com/patrickmn/go-cache v2.1.0+incompatible github.com/pion/dtls/v2 v2.1.5 github.com/pires/go-proxyproto v0.6.2 github.com/prometheus/client_golang v1.12.1 diff --git a/go.sum b/go.sum index 22c6e01..360f369 100644 --- a/go.sum +++ b/go.sum @@ -280,6 +280,8 @@ github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE= github.com/onsi/ginkgo/v2 v2.2.0 h1:3ZNA3L1c5FYDFTTxbFeVGGD8jYvjYauHD30YgLxVsNI= github.com/onsi/ginkgo/v2 v2.2.0/go.mod h1:MEH45j8TBi6u9BMogfbp0stKC5cdGjumZj5Y7AG4VIk= github.com/onsi/gomega v1.20.1 h1:PA/3qinGoukvymdIDV8pii6tiZgC8kbmJO6Z5+b002Q= +github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc= +github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ= github.com/pelletier/go-toml v1.9.4 h1:tjENF6MfZAg8e4ZmZTeWaWiT2vXtsoO6+iuOjFhECwM= github.com/pelletier/go-toml v1.9.4/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c= github.com/pion/dtls/v2 v2.1.5 h1:jlh2vtIyUBShchoTDqpCCqiYCyRFJ/lvf/gQ8TALs+c= diff --git a/limiter/traffic/generator.go b/limiter/traffic/generator.go index 4eaff07..2e17835 100644 --- a/limiter/traffic/generator.go +++ b/limiter/traffic/generator.go @@ -4,58 +4,28 @@ import ( limiter "github.com/go-gost/core/limiter/traffic" ) -type TrafficLimitGenerator interface { - In() limiter.Limiter - Out() limiter.Limiter -} - -type trafficLimitGenerator struct { +type limitGenerator struct { in int out int } -func NewTrafficLimitGenerator(in, out int) TrafficLimitGenerator { - return &trafficLimitGenerator{ +func newLimitGenerator(in, out int) *limitGenerator { + return &limitGenerator{ in: in, out: out, } } -func (p *trafficLimitGenerator) In() limiter.Limiter { +func (p *limitGenerator) In() limiter.Limiter { if p == nil || p.in <= 0 { return nil } return NewLimiter(p.in) } -func (p *trafficLimitGenerator) Out() limiter.Limiter { +func (p *limitGenerator) Out() limiter.Limiter { if p == nil || p.out <= 0 { return nil } return NewLimiter(p.out) } - -type trafficLimitSingleGenerator struct { - in limiter.Limiter - out limiter.Limiter -} - -func NewTrafficLimitSingleGenerator(in, out int) TrafficLimitGenerator { - p := &trafficLimitSingleGenerator{} - if in > 0 { - p.in = NewLimiter(in) - } - if out > 0 { - p.out = NewLimiter(out) - } - - return p -} - -func (p *trafficLimitSingleGenerator) In() limiter.Limiter { - return p.in -} - -func (p *trafficLimitSingleGenerator) Out() limiter.Limiter { - return p.out -} diff --git a/limiter/traffic/limiter.go b/limiter/traffic/limiter.go index 7cbcdd8..0a100d8 100644 --- a/limiter/traffic/limiter.go +++ b/limiter/traffic/limiter.go @@ -2,6 +2,9 @@ package traffic import ( "context" + "fmt" + "sort" + "strconv" limiter "github.com/go-gost/core/limiter/traffic" "golang.org/x/time/rate" @@ -26,5 +29,48 @@ func (l *llimiter) Wait(ctx context.Context, n int) int { } func (l *llimiter) Limit() int { - return l.limiter.Burst() + 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) } diff --git a/limiter/traffic/traffic.go b/limiter/traffic/traffic.go index 26309d8..95c29a7 100644 --- a/limiter/traffic/traffic.go +++ b/limiter/traffic/traffic.go @@ -5,7 +5,6 @@ import ( "context" "io" "net" - "sort" "strings" "sync" "time" @@ -14,6 +13,7 @@ import ( limiter "github.com/go-gost/core/limiter/traffic" "github.com/go-gost/core/logger" "github.com/go-gost/x/internal/loader" + "github.com/patrickmn/go-cache" "github.com/yl2chen/cidranger" ) @@ -22,33 +22,10 @@ const ( ConnLimitKey = "$$" ) -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() -} +const ( + defaultExpiration = 15 * time.Second + cleanupInterval = 30 * time.Second +) type options struct { limits []string @@ -97,14 +74,21 @@ func LoggerOption(logger logger.Logger) Option { } } +type limitValue struct { + in int + out int +} + type trafficLimiter struct { - limits map[string]TrafficLimitGenerator - cidrLimits cidranger.Ranger - inLimits map[string]limiter.Limiter - outLimits map[string]limiter.Limiter - mu sync.Mutex - cancelFunc context.CancelFunc - options options + generators sync.Map + cidrGenerators cidranger.Ranger + connInLimits *cache.Cache + connOutLimits *cache.Cache + inLimits *cache.Cache + outLimits *cache.Cache + mu sync.RWMutex + cancelFunc context.CancelFunc + options options } func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter { @@ -115,12 +99,13 @@ func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter { ctx, cancel := context.WithCancel(context.TODO()) lim := &trafficLimiter{ - limits: make(map[string]TrafficLimitGenerator), - cidrLimits: cidranger.NewPCTrieRanger(), - inLimits: make(map[string]limiter.Limiter), - outLimits: make(map[string]limiter.Limiter), - options: options, - cancelFunc: cancel, + cidrGenerators: cidranger.NewPCTrieRanger(), + connInLimits: cache.New(defaultExpiration, cleanupInterval), + connOutLimits: cache.New(defaultExpiration, cleanupInterval), + inLimits: cache.New(defaultExpiration, cleanupInterval), + outLimits: cache.New(defaultExpiration, cleanupInterval), + options: options, + cancelFunc: cancel, } if err := lim.reload(ctx); err != nil { @@ -132,45 +117,55 @@ func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter { return lim } +// In obtains a traffic input limiter based on key. +// The key should be client connection address. func (l *trafficLimiter) In(key string) limiter.Limiter { - l.mu.Lock() - defer l.mu.Unlock() - var lims []limiter.Limiter - if p := l.limits[ConnLimitKey]; p != nil { - if lim := p.In(); lim != nil { - lims = append(lims, lim) - } + // service level limiter + if lim, ok := l.inLimits.Get(GlobalLimitKey); ok && lim != nil { + lims = append(lims, lim.(limiter.Limiter)) } - if p := l.limits[GlobalLimitKey]; p != nil { - if lim := p.In(); lim != nil { - lims = append(lims, lim) + + // connection level limiter + if lim, ok := l.connInLimits.Get(key); ok { + if lim != nil { + // cached connection level limiter + lims = append(lims, lim.(limiter.Limiter)) + // reset expiration + l.connInLimits.Set(key, lim, defaultExpiration) + } + } else { + // generate a new connection level limiter and cache it + if v, ok := l.generators.Load(ConnLimitKey); ok && v != nil { + lim := v.(*limitGenerator).In() + if lim != nil { + lims = append(lims, lim) + l.connInLimits.Set(key, lim, defaultExpiration) + } } } - // IP limiter - if lim, ok := l.inLimits[key]; ok { + host, _, _ := net.SplitHostPort(key) + // IP level limiter + if lim, ok := l.inLimits.Get(host); ok { + // cached IP limiter if lim != nil { - lims = append(lims, lim) + lims = append(lims, lim.(limiter.Limiter)) } } else { - if ip := net.ParseIP(key); ip != nil { - if p := l.limits[key]; p != nil { - if lim = p.In(); lim != nil { + l.mu.RLock() + ranger := l.cidrGenerators + l.mu.RUnlock() + + // CIDR level limiter + if p, _ := ranger.ContainingNetworks(net.ParseIP(host)); len(p) > 0 { + if v, _ := p[0].(*cidrLimitEntry); v != nil { + if lim := v.generator.In(); lim != nil { lims = append(lims, lim) + l.inLimits.Set(host, lim, cache.NoExpiration) } } - if lim == nil { - if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 { - if v, _ := p[0].(*cidrLimitEntry); v != nil { - if lim = v.limit.In(); lim != nil { - lims = append(lims, lim) - } - } - } - } - l.inLimits[key] = lim } } @@ -180,51 +175,61 @@ func (l *trafficLimiter) In(key string) limiter.Limiter { } if lim != nil && l.options.logger != nil { - l.options.logger.Debugf("input limit for %s: %d", key, lim.Limit()) + l.options.logger.Debugf("input limit for %s: %s", key, lim) } return lim } +// Out obtains a traffic output limiter based on key. +// The key should be client connection address. func (l *trafficLimiter) Out(key string) limiter.Limiter { - l.mu.Lock() - defer l.mu.Unlock() - var lims []limiter.Limiter - if p := l.limits[ConnLimitKey]; p != nil { - if lim := p.Out(); lim != nil { - lims = append(lims, lim) - } + // service level limiter + if lim, ok := l.outLimits.Get(GlobalLimitKey); ok && lim != nil { + lims = append(lims, lim.(limiter.Limiter)) } - if p := l.limits[GlobalLimitKey]; p != nil { - if lim := p.Out(); lim != nil { - lims = append(lims, lim) + + // connection level limiter + if lim, ok := l.connOutLimits.Get(key); ok { + if lim != nil { + // cached connection level limiter + lims = append(lims, lim.(limiter.Limiter)) + // reset expiration + l.connOutLimits.Set(key, lim, defaultExpiration) + } + } else { + // generate a new connection level limiter + if v, ok := l.generators.Load(ConnLimitKey); ok && v != nil { + lim := v.(*limitGenerator).Out() + if lim != nil { + lims = append(lims, lim) + l.connOutLimits.Set(key, lim, defaultExpiration) + } } } - // IP limiter - if lim, ok := l.outLimits[key]; ok { + host, _, _ := net.SplitHostPort(key) + // IP level limiter + if lim, ok := l.outLimits.Get(host); ok { if lim != nil { - lims = append(lims, lim) + // cached IP level limiter + lims = append(lims, lim.(limiter.Limiter)) } } else { - if ip := net.ParseIP(key); ip != nil { - if p := l.limits[key]; p != nil { - if lim = p.Out(); lim != nil { + l.mu.RLock() + ranger := l.cidrGenerators + l.mu.RUnlock() + + // CIDR level limiter + if p, _ := ranger.ContainingNetworks(net.ParseIP(host)); len(p) > 0 { + if v, _ := p[0].(*cidrLimitEntry); v != nil { + if lim := v.generator.Out(); lim != nil { lims = append(lims, lim) + l.outLimits.Set(host, lim, cache.NoExpiration) } } - if lim == nil { - if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 { - if v, _ := p[0].(*cidrLimitEntry); v != nil { - if lim = v.limit.Out(); lim != nil { - lims = append(lims, lim) - } - } - } - } - l.outLimits[key] = lim } } @@ -234,7 +239,7 @@ func (l *trafficLimiter) Out(key string) limiter.Limiter { } if lim != nil && l.options.logger != nil { - l.options.logger.Debugf("output limit for %s: %d", key, lim.Limit()) + l.options.logger.Debugf("output limit for %s: %s", key, lim) } return lim @@ -262,36 +267,160 @@ func (l *trafficLimiter) periodReload(ctx context.Context) error { } func (l *trafficLimiter) reload(ctx context.Context) error { - v, err := l.load(ctx) + values, err := l.load(ctx) if err != nil { return err } - lines := append(l.options.limits, v...) - - limits := make(map[string]TrafficLimitGenerator) - cidrLimits := cidranger.NewPCTrieRanger() - - for _, s := range lines { - key, in, out := l.parseLimit(s) - if key == "" { - continue - } - switch key { - case GlobalLimitKey: - limits[key] = NewTrafficLimitSingleGenerator(in, out) - case ConnLimitKey: - limits[key] = NewTrafficLimitGenerator(in, out) - default: - if ip := net.ParseIP(key); ip != nil { - limits[key] = NewTrafficLimitSingleGenerator(in, out) - break + // service level limiter, never expired + { + value := values[GlobalLimitKey] + if v, _ := l.inLimits.Get(GlobalLimitKey); v != nil { + lim := v.(limiter.Limiter) + if value.in <= 0 { + l.inLimits.Delete(GlobalLimitKey) + } else { + lim.Set(value.in) } + } else { + if value.in > 0 { + l.inLimits.Set(GlobalLimitKey, NewLimiter(value.in), cache.NoExpiration) + } + } + + if v, _ := l.outLimits.Get(GlobalLimitKey); v != nil { + lim := v.(limiter.Limiter) + if value.out <= 0 { + l.outLimits.Delete(GlobalLimitKey) + } else { + lim.Set(value.out) + } + } else { + if value.out > 0 { + l.outLimits.Set(GlobalLimitKey, NewLimiter(value.out), cache.NoExpiration) + } + } + delete(values, GlobalLimitKey) + } + + // connection level limiters + { + value := values[ConnLimitKey] + + var in, out int + if v, _ := l.generators.Load(ConnLimitKey); v != nil { + in, out = v.(*limitGenerator).in, v.(*limitGenerator).out + } + l.generators.Store(ConnLimitKey, newLimitGenerator(value.in, value.out)) + + if value.in <= 0 { + l.connInLimits.Flush() + } else { + if in != value.in { + for _, item := range l.connInLimits.Items() { + if v := item.Object; v != nil { + v.(limiter.Limiter).Set(in) + } + } + } + } + + if value.out <= 0 { + l.connOutLimits.Flush() + } else { + if out != value.out { + for _, item := range l.connOutLimits.Items() { + if v := item.Object; v != nil { + v.(limiter.Limiter).Set(out) + } + } + } + } + delete(values, ConnLimitKey) + } + + cidrGenerators := cidranger.NewPCTrieRanger() + // IP/CIDR level limiters + { + // snapshot of the current limiters + inLimits := l.inLimits.Items() + outLimits := l.outLimits.Items() + + delete(inLimits, GlobalLimitKey) + delete(outLimits, GlobalLimitKey) + + for key, value := range values { if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil { - cidrLimits.Insert(&cidrLimitEntry{ - ipNet: *ipNet, - limit: NewTrafficLimitGenerator(in, out), + cidrGenerators.Insert(&cidrLimitEntry{ + ipNet: *ipNet, + generator: newLimitGenerator(value.in, value.out), }) + continue + } + + if v, _ := l.inLimits.Get(key); v != nil { + lim := v.(limiter.Limiter) + if value.in <= 0 { + l.inLimits.Delete(key) + } else { + lim.Set(value.in) + } + delete(inLimits, key) + } else { + if value.in > 0 { + l.inLimits.Set(key, NewLimiter(value.in), cache.NoExpiration) + } + } + + if v, _ := l.outLimits.Get(key); v != nil { + lim := v.(limiter.Limiter) + if value.out <= 0 { + l.outLimits.Delete(key) + } else { + lim.Set(value.out) + } + delete(outLimits, key) + } else { + if value.out > 0 { + l.outLimits.Set(key, NewLimiter(value.out), cache.NoExpiration) + } + } + } + + // check the CIDR for remain limiters, clean the unmatched ones. + for k, v := range inLimits { + if p, _ := cidrGenerators.ContainingNetworks(net.ParseIP(k)); len(p) > 0 { + if le, _ := p[0].(*cidrLimitEntry); le != nil { + in := le.generator.in + if in <= 0 { + l.inLimits.Delete(k) + continue + } + lim := v.Object.(limiter.Limiter) + if lim.Limit() != in { + lim.Set(in) + } + } + } else { + l.inLimits.Delete(k) + } + } + for k, v := range outLimits { + if p, _ := cidrGenerators.ContainingNetworks(net.ParseIP(k)); len(p) > 0 { + if le, _ := p[0].(*cidrLimitEntry); le != nil { + out := le.generator.out + if out <= 0 { + l.outLimits.Delete(k) + continue + } + lim := v.Object.(limiter.Limiter) + if lim.Limit() != out { + lim.Set(out) + } + delete(outLimits, k) + } + } else { + l.outLimits.Delete(k) } } } @@ -299,15 +428,22 @@ func (l *trafficLimiter) reload(ctx context.Context) error { l.mu.Lock() defer l.mu.Unlock() - l.limits = limits - l.cidrLimits = cidrLimits - l.inLimits = make(map[string]limiter.Limiter) - l.outLimits = make(map[string]limiter.Limiter) + l.cidrGenerators = cidrGenerators return nil } -func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error) { +func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue, err error) { + values = make(map[string]limitValue) + + for _, v := range l.options.limits { + key, in, out := l.parseLimit(v) + if key == "" { + continue + } + values[key] = limitValue{in: in, out: out} + } + if l.options.fileLoader != nil { if lister, ok := l.options.fileLoader.(loader.Lister); ok { list, er := lister.List(ctx) @@ -315,17 +451,24 @@ func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error l.options.logger.Warnf("file loader: %v", er) } for _, s := range list { - if line := l.parseLine(s); line != "" { - patterns = append(patterns, line) + key, in, out := l.parseLimit(l.parseLine(s)) + if key == "" { + continue } + values[key] = limitValue{in: in, out: out} } } else { r, er := l.options.fileLoader.Load(ctx) if er != nil { l.options.logger.Warnf("file loader: %v", er) } - if v, _ := l.parsePatterns(r); v != nil { - patterns = append(patterns, v...) + patterns, _ := l.parsePatterns(r) + for _, s := range patterns { + key, in, out := l.parseLimit(l.parseLine(s)) + if key == "" { + continue + } + values[key] = limitValue{in: in, out: out} } } } @@ -335,14 +478,25 @@ func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error if er != nil { l.options.logger.Warnf("redis loader: %v", er) } - patterns = append(patterns, list...) + for _, s := range list { + key, in, out := l.parseLimit(l.parseLine(s)) + if key == "" { + continue + } + values[key] = limitValue{in: in, out: out} + } } else { r, er := l.options.redisLoader.Load(ctx) if er != nil { l.options.logger.Warnf("redis loader: %v", er) } - if v, _ := l.parsePatterns(r); v != nil { - patterns = append(patterns, v...) + patterns, _ := l.parsePatterns(r) + for _, s := range patterns { + key, in, out := l.parseLimit(l.parseLine(s)) + if key == "" { + continue + } + values[key] = limitValue{in: in, out: out} } } } @@ -351,12 +505,17 @@ func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error if er != nil { l.options.logger.Warnf("http loader: %v", er) } - if v, _ := l.parsePatterns(r); v != nil { - patterns = append(patterns, v...) + patterns, _ := l.parsePatterns(r) + for _, s := range patterns { + key, in, out := l.parseLimit(l.parseLine(s)) + if key == "" { + continue + } + values[key] = limitValue{in: in, out: out} } } - l.options.logger.Debugf("load items %d", len(patterns)) + l.options.logger.Debugf("load items %d", len(values)) return } @@ -386,6 +545,10 @@ func (l *trafficLimiter) parseLine(s string) string { func (l *trafficLimiter) parseLimit(s string) (key string, in, out int) { s = strings.Replace(s, "\t", " ", -1) s = strings.TrimSpace(s) + if s == "" { + return + } + var ss []string for _, v := range strings.Split(s, " ") { if v != "" { @@ -421,8 +584,8 @@ func (l *trafficLimiter) Close() error { } type cidrLimitEntry struct { - ipNet net.IPNet - limit TrafficLimitGenerator + ipNet net.IPNet + generator *limitGenerator } func (p *cidrLimitEntry) Network() net.IPNet { diff --git a/limiter/traffic/wrapper/conn.go b/limiter/traffic/wrapper/conn.go index d75a4f7..3842299 100644 --- a/limiter/traffic/wrapper/conn.go +++ b/limiter/traffic/wrapper/conn.go @@ -6,40 +6,63 @@ import ( "errors" "io" "net" - "sync" "syscall" + "time" limiter "github.com/go-gost/core/limiter/traffic" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/udp" + "github.com/patrickmn/go-cache" ) var ( errUnsupport = errors.New("unsupported operation") ) -// serverConn is a server side Conn with metrics supported. +// serverConn is a server side Conn with traffic limiter supported. type serverConn struct { net.Conn rbuf bytes.Buffer + limiter limiter.TrafficLimiter limiterIn limiter.Limiter + expIn int64 limiterOut limiter.Limiter + expOut int64 } -func WrapConn(rlimiter limiter.TrafficLimiter, c net.Conn) net.Conn { - if rlimiter == nil { +func WrapConn(limiter limiter.TrafficLimiter, c net.Conn) net.Conn { + if limiter == nil { return c } - host, _, _ := net.SplitHostPort(c.RemoteAddr().String()) return &serverConn{ - Conn: c, - limiterIn: rlimiter.In(host), - limiterOut: rlimiter.Out(host), + Conn: c, + limiter: limiter, } } +func (c *serverConn) getInLimiter(addr net.Addr) limiter.Limiter { + now := time.Now().UnixNano() + // cache the limiter for 1s + if c.limiter != nil && time.Duration(now-c.expIn) > time.Second { + c.limiterIn = c.limiter.In(addr.String()) + c.expIn = now + } + return c.limiterIn +} + +func (c *serverConn) getOutLimiter(addr net.Addr) limiter.Limiter { + now := time.Now().UnixNano() + // cache the limiter for 1s + if c.limiter != nil && time.Duration(now-c.expOut) > time.Second { + c.limiterOut = c.limiter.Out(addr.String()) + c.expOut = now + } + return c.limiterOut +} + func (c *serverConn) Read(b []byte) (n int, err error) { - if c.limiterIn == nil { + limiter := c.getInLimiter(c.RemoteAddr()) + if limiter == nil { return c.Conn.Read(b) } @@ -48,7 +71,7 @@ func (c *serverConn) Read(b []byte) (n int, err error) { if c.rbuf.Len() < burst { burst = c.rbuf.Len() } - lim := c.limiterIn.Wait(context.Background(), burst) + lim := limiter.Wait(context.Background(), burst) return c.rbuf.Read(b[:lim]) } @@ -57,7 +80,7 @@ func (c *serverConn) Read(b []byte) (n int, err error) { return nn, err } - n = c.limiterIn.Wait(context.Background(), nn) + n = limiter.Wait(context.Background(), nn) if n < nn { if _, err = c.rbuf.Write(b[n:nn]); err != nil { return 0, err @@ -68,13 +91,14 @@ func (c *serverConn) Read(b []byte) (n int, err error) { } func (c *serverConn) Write(b []byte) (n int, err error) { - if c.limiterOut == nil { + limiter := c.getOutLimiter(c.RemoteAddr()) + if limiter == nil { return c.Conn.Write(b) } nn := 0 for len(b) > 0 { - nn, err = c.Conn.Write(b[:c.limiterOut.Wait(context.Background(), len(b))]) + nn, err = c.Conn.Write(b[:limiter.Wait(context.Background(), len(b))]) n += nn if err != nil { return @@ -97,10 +121,8 @@ func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) { type packetConn struct { net.PacketConn limiter limiter.TrafficLimiter - inLimits map[string]limiter.Limiter - inMux sync.RWMutex - outLimits map[string]limiter.Limiter - outMux sync.RWMutex + inLimits *cache.Cache + outLimits *cache.Cache } func WrapPacketConn(lim limiter.TrafficLimiter, pc net.PacketConn) net.PacketConn { @@ -110,8 +132,8 @@ func WrapPacketConn(lim limiter.TrafficLimiter, pc net.PacketConn) net.PacketCon return &packetConn{ PacketConn: pc, limiter: lim, - inLimits: make(map[string]limiter.Limiter), - outLimits: make(map[string]limiter.Limiter), + inLimits: cache.New(time.Second, 10*time.Second), + outLimits: cache.New(time.Second, 10*time.Second), } } @@ -120,24 +142,21 @@ func (c *packetConn) getInLimiter(addr net.Addr) limiter.Limiter { return nil } - lim, ok := func() (limiter.Limiter, bool) { - c.inMux.RLock() - defer c.inMux.RUnlock() - - lim, ok := c.inLimits[addr.String()] - return lim, ok + lim, ok := func() (lim limiter.Limiter, ok bool) { + v, ok := c.inLimits.Get(addr.String()) + if ok { + if v != nil { + lim = v.(limiter.Limiter) + } + } + return }() if ok { return lim } - host, _, _ := net.SplitHostPort(addr.String()) - lim = c.limiter.In(host) - - c.inMux.Lock() - defer c.inMux.Unlock() - - c.inLimits[addr.String()] = lim + lim = c.limiter.In(addr.String()) + c.inLimits.Set(addr.String(), lim, 0) return lim } @@ -147,24 +166,21 @@ func (c *packetConn) getOutLimiter(addr net.Addr) limiter.Limiter { return nil } - lim, ok := func() (limiter.Limiter, bool) { - c.outMux.RLock() - defer c.outMux.RUnlock() - - lim, ok := c.outLimits[addr.String()] - return lim, ok + lim, ok := func() (lim limiter.Limiter, ok bool) { + v, ok := c.outLimits.Get(addr.String()) + if ok { + if v != nil { + lim = v.(limiter.Limiter) + } + } + return }() if ok { return lim } - host, _, _ := net.SplitHostPort(addr.String()) - lim = c.limiter.Out(host) - - c.outMux.Lock() - defer c.outMux.Unlock() - - c.outLimits[addr.String()] = lim + lim = c.limiter.Out(addr.String()) + c.outLimits.Set(addr.String(), lim, 0) return lim } @@ -204,16 +220,16 @@ func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { type udpConn struct { net.PacketConn limiter limiter.TrafficLimiter - inLimits map[string]limiter.Limiter - inMux sync.RWMutex - outLimits map[string]limiter.Limiter - outMux sync.RWMutex + inLimits *cache.Cache + outLimits *cache.Cache } func WrapUDPConn(limiter limiter.TrafficLimiter, pc net.PacketConn) udp.Conn { return &udpConn{ PacketConn: pc, limiter: limiter, + inLimits: cache.New(time.Second, 10*time.Second), + outLimits: cache.New(time.Second, 10*time.Second), } } @@ -222,24 +238,21 @@ func (c *udpConn) getInLimiter(addr net.Addr) limiter.Limiter { return nil } - lim, ok := func() (limiter.Limiter, bool) { - c.inMux.RLock() - defer c.inMux.RUnlock() - - lim, ok := c.inLimits[addr.String()] - return lim, ok + lim, ok := func() (lim limiter.Limiter, ok bool) { + v, ok := c.inLimits.Get(addr.String()) + if ok { + if v != nil { + lim = v.(limiter.Limiter) + } + } + return }() if ok { return lim } - host, _, _ := net.SplitHostPort(addr.String()) - lim = c.limiter.In(host) - - c.inMux.Lock() - defer c.inMux.Unlock() - - c.inLimits[addr.String()] = lim + lim = c.limiter.In(addr.String()) + c.inLimits.Set(addr.String(), lim, 0) return lim } @@ -249,24 +262,21 @@ func (c *udpConn) getOutLimiter(addr net.Addr) limiter.Limiter { return nil } - lim, ok := func() (limiter.Limiter, bool) { - c.outMux.RLock() - defer c.outMux.RUnlock() - - lim, ok := c.outLimits[addr.String()] - return lim, ok + lim, ok := func() (lim limiter.Limiter, ok bool) { + v, ok := c.outLimits.Get(addr.String()) + if ok { + if v != nil { + lim = v.(limiter.Limiter) + } + } + return }() if ok { return lim } - host, _, _ := net.SplitHostPort(addr.String()) - lim = c.limiter.Out(host) - - c.outMux.Lock() - defer c.outMux.Unlock() - - c.outLimits[addr.String()] = lim + lim = c.limiter.Out(addr.String()) + c.outLimits.Set(addr.String(), lim, 0) return lim }