go-gost/x
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(limiter/traffic): add burst support, dropped-packet counters, and cache fixes

Browse Source 
- 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
This commit is contained in:
ginuerzh
2026-05-24 21:52:02 +08:00
parent a9d94e7009
commit 64f9a048f8
9 changed files with 207 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
limiter/traffic/cache/limiter.go
Vendored
+6
View File
@@ -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) {
limiter/traffic/generator.go
+13 -5
View File
@@ -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)
}
limiter/traffic/generator_test.go
+29 -4
View File
@@ -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())
}
}
limiter/traffic/limiter.go
+12 -1
View File
@@ -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()
limiter/traffic/limiter_test.go
+34
View File
@@ -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)
limiter/traffic/traffic.go
+39 -25
View File
@@ -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
}
limiter/traffic/traffic_test.go
+4 -1
View File
@@ -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)
limiter/traffic/wrapper/conn.go
+33 -5
View File
@@ -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
}
}
limiter/traffic/wrapper/conn_test.go
+37 -1
View File
@@ -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
}