diff --git a/limiter/traffic/cache/limiter.go b/limiter/traffic/cache/limiter.go index 70c497de..e6c1e3c2 100644 --- a/limiter/traffic/cache/limiter.go +++ b/limiter/traffic/cache/limiter.go @@ -115,6 +115,9 @@ func (p *cachedTrafficLimiter) In(ctx context.Context, key string, opts ...limit limNew := p.limiter.In(ctx, key, opts...) if limNew == nil { + if lim == nil { + return nil + } limNew = lim } if item == nil || !p.equal(lim, limNew) { @@ -149,6 +152,9 @@ func (p *cachedTrafficLimiter) Out(ctx context.Context, key string, opts ...limi limNew := p.limiter.Out(ctx, key, opts...) if limNew == nil { + if lim == nil { + return nil + } limNew = lim } if item == nil || !p.equal(lim, limNew) { diff --git a/limiter/traffic/generator.go b/limiter/traffic/generator.go index 2e178356..298046f2 100644 --- a/limiter/traffic/generator.go +++ b/limiter/traffic/generator.go @@ -5,14 +5,16 @@ import ( ) type limitGenerator struct { - in int - out int + in int + out int + burst int } -func newLimitGenerator(in, out int) *limitGenerator { +func newLimitGenerator(in, out, burst int) *limitGenerator { return &limitGenerator{ - in: in, - out: out, + in: in, + out: out, + burst: burst, } } @@ -20,6 +22,9 @@ func (p *limitGenerator) In() limiter.Limiter { if p == nil || p.in <= 0 { return nil } + if p.burst > 0 { + return NewLimiterWithBurst(p.in, p.burst) + } return NewLimiter(p.in) } @@ -27,5 +32,8 @@ func (p *limitGenerator) Out() limiter.Limiter { if p == nil || p.out <= 0 { return nil } + if p.burst > 0 { + return NewLimiterWithBurst(p.out, p.burst) + } return NewLimiter(p.out) } diff --git a/limiter/traffic/generator_test.go b/limiter/traffic/generator_test.go index 06b0d072..3a85f9a6 100644 --- a/limiter/traffic/generator_test.go +++ b/limiter/traffic/generator_test.go @@ -1,11 +1,12 @@ package traffic import ( + "context" "testing" ) func TestLimitGenerator_In(t *testing.T) { - g := newLimitGenerator(100, 0) + g := newLimitGenerator(100, 0, 0) lim := g.In() if lim == nil { t.Fatal("In() should not be nil for positive in") @@ -22,14 +23,14 @@ func TestLimitGenerator_In(t *testing.T) { } func TestLimitGenerator_In_Zero(t *testing.T) { - g := newLimitGenerator(0, 200) + g := newLimitGenerator(0, 200, 0) if lim := g.In(); lim != nil { t.Fatal("In() should be nil for zero in") } } func TestLimitGenerator_Out(t *testing.T) { - g := newLimitGenerator(0, 200) + g := newLimitGenerator(0, 200, 0) lim := g.Out() if lim == nil { t.Fatal("Out() should not be nil for positive out") @@ -46,7 +47,7 @@ func TestLimitGenerator_Out(t *testing.T) { } func TestLimitGenerator_Out_Zero(t *testing.T) { - g := newLimitGenerator(100, 0) + g := newLimitGenerator(100, 0, 0) if lim := g.Out(); lim != nil { t.Fatal("Out() should be nil for zero out") } @@ -61,3 +62,27 @@ func TestLimitGenerator_NilReceiver(t *testing.T) { t.Fatal("nil receiver Out() should return nil") } } + +func TestLimitGenerator_Burst(t *testing.T) { + // Burst specified: should use NewLimiterWithBurst. + g := newLimitGenerator(100, 200, 500) + limIn := g.In() + if limIn == nil { + t.Fatal("In() should not be nil") + } + if limIn.Limit() != 100 { + t.Fatalf("expected limit 100, got %d", limIn.Limit()) + } + // With burst=500, Wait for 500 should succeed immediately. + if n := limIn.Wait(context.Background(), 500); n != 500 { + t.Fatalf("expected 500 with burst, got %d", n) + } + + limOut := g.Out() + if limOut == nil { + t.Fatal("Out() should not be nil") + } + if limOut.Limit() != 200 { + t.Fatalf("expected limit 200, got %d", limOut.Limit()) + } +} diff --git a/limiter/traffic/limiter.go b/limiter/traffic/limiter.go index eb9ae12f..e608058c 100644 --- a/limiter/traffic/limiter.go +++ b/limiter/traffic/limiter.go @@ -15,13 +15,24 @@ type llimiter struct { } // NewLimiter creates a rate-based traffic Limiter with the given rate -// (bytes per second) and burst size. +// (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() diff --git a/limiter/traffic/limiter_test.go b/limiter/traffic/limiter_test.go index d5cb85e5..53de7052 100644 --- a/limiter/traffic/limiter_test.go +++ b/limiter/traffic/limiter_test.go @@ -17,6 +17,40 @@ func TestNewLimiter(t *testing.T) { } } +func TestNewLimiterWithBurst(t *testing.T) { + l := NewLimiterWithBurst(100, 500) + if l == nil { + t.Fatal("NewLimiterWithBurst should not be nil") + } + if l.Limit() != 100 { + t.Fatalf("expected limit 100, got %d", l.Limit()) + } + // With burst=500, Wait for 500 should succeed immediately. + if n := l.Wait(context.Background(), 500); n != 500 { + t.Fatalf("expected 500 with burst, got %d", n) + } + // With burst=500, Wait for 501 clamps to burst. + if n := l.Wait(context.Background(), 501); n != 500 { + t.Fatalf("expected burst 500, got %d", n) + } +} + +func TestNewLimiterWithBurst_Defaults(t *testing.T) { + // burst <= 0 defaults to rate. + l := NewLimiterWithBurst(100, 0) + if l.Limit() != 100 { + t.Fatalf("expected limit 100, got %d", l.Limit()) + } + if n := l.Wait(context.Background(), 100); n != 100 { + t.Fatalf("expected burst=rate=100, got %d", n) + } + + l2 := NewLimiterWithBurst(50, -1) + if n := l2.Wait(context.Background(), 50); n != 50 { + t.Fatalf("expected burst=rate=50, got %d", n) + } +} + func TestLimiter_Wait(t *testing.T) { // High burst so Wait does not block. l := NewLimiter(1000000) diff --git a/limiter/traffic/traffic.go b/limiter/traffic/traffic.go index 23504872..b70c2ed4 100644 --- a/limiter/traffic/traffic.go +++ b/limiter/traffic/traffic.go @@ -92,8 +92,9 @@ func LoggerOption(logger logger.Logger) Option { } type limitValue struct { - in int - out int + in int + out int + burst int } type trafficLimiter struct { @@ -193,6 +194,8 @@ func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Opt // cached IP limiter if lim != nil { lims = append(lims, lim.(traffic.Limiter)) + // reset expiration to prevent unbounded cache growth + l.inLimits.Set(host, lim, defaultExpiration) } } else { l.mu.RLock() @@ -204,14 +207,16 @@ func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Opt 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) + l.inLimits.Set(host, lim, defaultExpiration) } } } } var lim traffic.Limiter - if len(lims) > 0 { + if len(lims) == 1 { + lim = lims[0] + } else if len(lims) > 1 { lim = newLimiterGroup(lims...) } @@ -275,6 +280,8 @@ func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Op if lim != nil { // cached IP level limiter lims = append(lims, lim.(traffic.Limiter)) + // reset expiration to prevent unbounded cache growth + l.outLimits.Set(host, lim, defaultExpiration) } } else { l.mu.RLock() @@ -286,14 +293,16 @@ func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Op 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) + l.outLimits.Set(host, lim, defaultExpiration) } } } } var lim traffic.Limiter - if len(lims) > 0 { + if len(lims) == 1 { + lim = lims[0] + } else if len(lims) > 1 { lim = newLimiterGroup(lims...) } @@ -351,7 +360,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error { } } else { if value.in > 0 { - l.inLimits.Set(ServiceLimitKey, NewLimiter(value.in), cache.NoExpiration) + l.inLimits.Set(ServiceLimitKey, NewLimiterWithBurst(value.in, value.burst), cache.NoExpiration) } } @@ -364,7 +373,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error { } } else { if value.out > 0 { - l.outLimits.Set(ServiceLimitKey, NewLimiter(value.out), cache.NoExpiration) + l.outLimits.Set(ServiceLimitKey, NewLimiterWithBurst(value.out, value.burst), cache.NoExpiration) } } delete(values, ServiceLimitKey) @@ -378,7 +387,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error { 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)) + l.generators.Store(ConnLimitKey, newLimitGenerator(value.in, value.out, value.burst)) if value.in <= 0 { l.connInLimits.Flush() @@ -420,7 +429,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error { if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil { cidrGenerators.Insert(&cidrLimitEntry{ ipNet: *ipNet, - generator: newLimitGenerator(value.in, value.out), + generator: newLimitGenerator(value.in, value.out, value.burst), }) continue } @@ -435,7 +444,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error { delete(inLimits, key) } else { if value.in > 0 { - l.inLimits.Set(key, NewLimiter(value.in), cache.NoExpiration) + l.inLimits.Set(key, NewLimiterWithBurst(value.in, value.burst), defaultExpiration) } } @@ -449,7 +458,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error { delete(outLimits, key) } else { if value.out > 0 { - l.outLimits.Set(key, NewLimiter(value.out), cache.NoExpiration) + l.outLimits.Set(key, NewLimiterWithBurst(value.out, value.burst), defaultExpiration) } } } @@ -504,11 +513,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue values = make(map[string]limitValue) for _, v := range l.options.limits { - key, in, out := l.parseLimit(v) + key, in, out, burst := l.parseLimit(v) if key == "" { continue } - values[key] = limitValue{in: in, out: out} + values[key] = limitValue{in: in, out: out, burst: burst} } if l.options.fileLoader != nil { @@ -518,11 +527,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue l.logger.Warnf("file loader: %v", er) } for _, s := range list { - key, in, out := l.parseLimit(l.parseLine(s)) + key, in, out, burst := l.parseLimit(l.parseLine(s)) if key == "" { continue } - values[key] = limitValue{in: in, out: out} + values[key] = limitValue{in: in, out: out, burst: burst} } } else { r, er := l.options.fileLoader.Load(ctx) @@ -531,11 +540,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue } patterns, _ := l.parsePatterns(r) for _, s := range patterns { - key, in, out := l.parseLimit(l.parseLine(s)) + key, in, out, burst := l.parseLimit(l.parseLine(s)) if key == "" { continue } - values[key] = limitValue{in: in, out: out} + values[key] = limitValue{in: in, out: out, burst: burst} } } } @@ -546,11 +555,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue l.logger.Warnf("redis loader: %v", er) } for _, s := range list { - key, in, out := l.parseLimit(l.parseLine(s)) + key, in, out, burst := l.parseLimit(l.parseLine(s)) if key == "" { continue } - values[key] = limitValue{in: in, out: out} + values[key] = limitValue{in: in, out: out, burst: burst} } } else { r, er := l.options.redisLoader.Load(ctx) @@ -559,11 +568,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue } patterns, _ := l.parsePatterns(r) for _, s := range patterns { - key, in, out := l.parseLimit(l.parseLine(s)) + key, in, out, burst := l.parseLimit(l.parseLine(s)) if key == "" { continue } - values[key] = limitValue{in: in, out: out} + values[key] = limitValue{in: in, out: out, burst: burst} } } } @@ -574,11 +583,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue } patterns, _ := l.parsePatterns(r) for _, s := range patterns { - key, in, out := l.parseLimit(l.parseLine(s)) + key, in, out, burst := l.parseLimit(l.parseLine(s)) if key == "" { continue } - values[key] = limitValue{in: in, out: out} + values[key] = limitValue{in: in, out: out, burst: burst} } } @@ -609,7 +618,7 @@ func (l *trafficLimiter) parseLine(s string) string { return strings.TrimSpace(s) } -func (l *trafficLimiter) parseLimit(s string) (key string, in, out int) { +func (l *trafficLimiter) parseLimit(s string) (key string, in, out, burst int) { s = strings.Replace(s, "\t", " ", -1) s = strings.TrimSpace(s) if s == "" { @@ -635,6 +644,11 @@ func (l *trafficLimiter) parseLimit(s string) (key string, in, out int) { out = int(v) } } + if len(ss) > 3 { + if v, _ := units.ParseBase2Bytes(ss[3]); v > 0 { + burst = int(v) + } + } return } diff --git a/limiter/traffic/traffic_test.go b/limiter/traffic/traffic_test.go index e96e0acd..f7194fb9 100644 --- a/limiter/traffic/traffic_test.go +++ b/limiter/traffic/traffic_test.go @@ -110,7 +110,10 @@ func TestParseLimit(t *testing.T) { {"key invalid_bytes", "key", 0, 0}, } for _, tt := range tests { - key, in, out := l.parseLimit(tt.input) + key, in, out, burst := l.parseLimit(tt.input) + if burst != 0 { + t.Errorf("parseLimit(%q) unexpected burst: %d", tt.input, burst) + } if key != tt.key || in != tt.in || out != tt.out { t.Errorf("parseLimit(%q) = (%q, %d, %d), want (%q, %d, %d)", tt.input, key, in, out, tt.key, tt.in, tt.out) diff --git a/limiter/traffic/wrapper/conn.go b/limiter/traffic/wrapper/conn.go index cc4d46b7..ee670787 100644 --- a/limiter/traffic/wrapper/conn.go +++ b/limiter/traffic/wrapper/conn.go @@ -8,6 +8,7 @@ import ( "errors" "io" "net" + "sync/atomic" "syscall" "github.com/go-gost/core/limiter" @@ -23,6 +24,16 @@ var ( errRateLimited = errors.New("rate limited") ) +// ErrRateLimited is returned when a write exceeds the rate limit. +var ErrRateLimited = errRateLimited + +// DroppedPacketCounter is an optional interface implemented by rate-limited +// packet connections that report the number of packets discarded due to rate +// limiting. +type DroppedPacketCounter interface { + DroppedPackets() int64 +} + // limitConn wraps a net.Conn with traffic rate limiting applied to reads and writes. type limitConn struct { net.Conn @@ -135,6 +146,12 @@ type packetConn struct { limiter traffic.TrafficLimiter opts []limiter.Option key string + dropped atomic.Int64 +} + +// DroppedPackets returns the number of packets discarded due to rate limiting. +func (c *packetConn) DroppedPackets() int64 { + return c.dropped.Load() } // WrapPacketConn wraps a net.PacketConn with traffic rate limiting. Packets @@ -165,6 +182,7 @@ func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { // discard when exceed the limit size. if limiter.Wait(context.Background(), n) < n { + c.dropped.Add(1) continue } @@ -176,7 +194,7 @@ func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { limiter := c.limiter.Out(context.Background(), c.key, c.opts...) if limiter != nil && limiter.Limit() > 0 && limiter.Wait(context.Background(), len(p)) < len(p) { - return 0, errRateLimited + return 0, ErrRateLimited } return c.PacketConn.WriteTo(p, addr) @@ -194,6 +212,12 @@ type udpConn struct { limiter traffic.TrafficLimiter opts []limiter.Option key string + dropped atomic.Int64 +} + +// DroppedPackets returns the number of packets discarded due to rate limiting. +func (c *udpConn) DroppedPackets() int64 { + return c.dropped.Load() } // WrapUDPConn wraps a net.PacketConn as a udp.Conn with traffic rate limiting. @@ -251,6 +275,7 @@ func (c *udpConn) Read(b []byte) (n int, err error) { // discard when exceed the limit size. if limiter.Wait(context.Background(), n) < n { + c.dropped.Add(1) continue } @@ -276,6 +301,7 @@ func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { // discard when exceed the limit size. if limiter.Wait(context.Background(), n) < n { + c.dropped.Add(1) continue } @@ -307,6 +333,7 @@ func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) { // discard when exceed the limit size. if limiter.Wait(context.Background(), n) < n { + c.dropped.Add(1) continue } @@ -338,6 +365,7 @@ func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAd // discard when exceed the limit size. if limiter.Wait(context.Background(), n) < n { + c.dropped.Add(1) continue } return @@ -355,7 +383,7 @@ func (c *udpConn) Write(p []byte) (n int, err error) { limiter := c.limiter.Out(context.Background(), c.key, c.opts...) if limiter != nil && limiter.Limit() > 0 && limiter.Wait(context.Background(), len(p)) < len(p) { - return 0, errRateLimited + return 0, ErrRateLimited } } @@ -368,7 +396,7 @@ func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { limiter := c.limiter.Out(context.Background(), c.key, c.opts...) if limiter != nil && limiter.Limit() > 0 && limiter.Wait(context.Background(), len(p)) < len(p) { - return 0, errRateLimited + return 0, ErrRateLimited } } @@ -387,7 +415,7 @@ func (c *udpConn) WriteToUDP(p []byte, addr *net.UDPAddr) (n int, err error) { limiter := c.limiter.Out(context.Background(), c.key, c.opts...) if limiter != nil && limiter.Limit() > 0 && limiter.Wait(context.Background(), len(p)) < len(p) { - return 0, errRateLimited + return 0, ErrRateLimited } } @@ -406,7 +434,7 @@ func (c *udpConn) WriteMsgUDP(p, oob []byte, addr *net.UDPAddr) (n, oobn int, er limiter := c.limiter.Out(context.Background(), c.key, c.opts...) if limiter != nil && limiter.Limit() > 0 && limiter.Wait(context.Background(), len(p)) < len(p) { - return 0, 0, errRateLimited + return 0, 0, ErrRateLimited } } diff --git a/limiter/traffic/wrapper/conn_test.go b/limiter/traffic/wrapper/conn_test.go index fb4e2f2e..291c8225 100644 --- a/limiter/traffic/wrapper/conn_test.go +++ b/limiter/traffic/wrapper/conn_test.go @@ -362,6 +362,35 @@ func TestWrapUDPConn(t *testing.T) { } } +func TestPacketConn_DroppedPackets(t *testing.T) { + pc := &mockPacketConn{writeBuf: make([]byte, 100)} + calls := 0 + pc.readFromFunc = func(p []byte) (int, net.Addr, error) { + calls++ + if calls == 1 { + return copy(p, "data-4"), nil, nil // 6 bytes, limiter allows 5 = drop + } + return copy(p, "ok"), nil, nil // 2 bytes, limiter allows 5 = pass + } + ml := &mockLimiter{ + limit: 100, + waitFunc: func(ctx context.Context, n int) int { + return 5 // allow 5 bytes; 6-byte packet is dropped, 2-byte passes + }, + } + tl := &mockTrafficLimiter{inLim: ml} + wrapped := WrapPacketConn(pc, tl, "test-key") + + _, _, err := wrapped.ReadFrom(make([]byte, 10)) + if err != nil { + t.Fatal(err) + } + counter := wrapped.(DroppedPacketCounter) + if n := counter.DroppedPackets(); n != 1 { + t.Fatalf("expected 1 dropped packet, got %d", n) + } +} + func TestUDPConn_Write_RateLimited(t *testing.T) { pc := &mockPacketConn{writeBuf: make([]byte, 100)} ml := &mockLimiter{ @@ -388,6 +417,9 @@ var ( _ xio.CloseWrite = (*limitConn)(nil) _ syscall.Conn = (*limitConn)(nil) _ ctxutil.Context = (*limitConn)(nil) + + _ DroppedPacketCounter = (*packetConn)(nil) + _ DroppedPacketCounter = (*udpConn)(nil) ) // --- helpers --- @@ -412,7 +444,8 @@ func (rw *bytesReadWriter) Write(b []byte) (int, error) { type mockPacketConn struct { net.PacketConn - writeBuf []byte + writeBuf []byte + readFromFunc func(p []byte) (int, net.Addr, error) } func (pc *mockPacketConn) Write(p []byte) (int, error) { @@ -426,6 +459,9 @@ func (pc *mockPacketConn) WriteTo(p []byte, addr net.Addr) (int, error) { } func (pc *mockPacketConn) ReadFrom(p []byte) (int, net.Addr, error) { + if pc.readFromFunc != nil { + return pc.readFromFunc(p) + } return copy(p, "test"), nil, nil }