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

- 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
2026-05-24 21:52:02 +08:00
parent a9d94e7009
commit 64f9a048f8
9 changed files with 207 additions and 42 deletions
@@ -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
 		}
 	}

@@ -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
 }