From fd17c3e18dd4426bf7968bdacefcec88a0e35bb4 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Sun, 14 Jun 2026 01:07:41 +0800
Subject: [PATCH] feat(observer/stats): track UDP connection counts in
 WrapUDPConn
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Mirror WrapConn behavior: WrapUDPConn now increments total and current
connection counts on creation (one accepted UDP client session == one
connection) and decrements the current count on Close. Close is guarded
by a mutex and a closed channel so it is safe to call multiple times —
the current-conns decrement fires exactly once and the monotonic total
is untouched.
---
 observer/stats/wrapper/conn.go      | 41 ++++++++++++++++++++++++-----
 observer/stats/wrapper/conn_test.go | 35 ++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/observer/stats/wrapper/conn.go b/observer/stats/wrapper/conn.go
index 3ccd88e9..b0385d0e 100644
--- a/observer/stats/wrapper/conn.go
+++ b/observer/stats/wrapper/conn.go
@@ -146,23 +146,50 @@ func (c *packetConn) Context() context.Context {
 
 type udpConn struct {
 	net.PacketConn
-	stats stats.Stats
+	stats  stats.Stats
+	closed chan struct{}
+	mu     sync.Mutex
 }
 
-// WrapUDPConn wraps a net.PacketConn as a udp.Conn to track input and output
-// bytes across all read/write methods (Read, Write, ReadFrom, WriteTo,
-// ReadFromUDP, ReadMsgUDP, WriteToUDP, WriteMsgUDP). If pc or stats is nil,
+// WrapUDPConn wraps a net.PacketConn as a udp.Conn to track connection and
+// traffic statistics. It increments total and current connection counts (one
+// accepted UDP client session == one connection, mirroring WrapConn) and
+// tracks bytes read and written across all read/write methods (Read, Write,
+// ReadFrom, WriteTo, ReadFromUDP, ReadMsgUDP, WriteToUDP, WriteMsgUDP). On
+// Close, it decrements the current connection count. If pc or pStats is nil,
 // it returns nil.
-func WrapUDPConn(pc net.PacketConn, stats stats.Stats) udp.Conn {
-	if pc == nil || stats == nil {
+func WrapUDPConn(pc net.PacketConn, pStats stats.Stats) udp.Conn {
+	if pc == nil || pStats == nil {
 		return nil
 	}
+
+	pStats.Add(stats.KindTotalConns, 1)
+	pStats.Add(stats.KindCurrentConns, 1)
+
 	return &udpConn{
 		PacketConn: pc,
-		stats:      stats,
+		stats:      pStats,
+		closed:     make(chan struct{}),
 	}
 }
 
+// Close decrements the current connection count once and closes the underlying
+// PacketConn. It is safe to call multiple times.
+func (c *udpConn) Close() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	select {
+	case <-c.closed:
+		return nil
+	default:
+		close(c.closed)
+	}
+
+	c.stats.Add(stats.KindCurrentConns, -1)
+	return c.PacketConn.Close()
+}
+
 func (c *udpConn) RemoteAddr() net.Addr {
 	if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok {
 		return nc.RemoteAddr()
diff --git a/observer/stats/wrapper/conn_test.go b/observer/stats/wrapper/conn_test.go
index 66acb63d..dd8e69e2 100644
--- a/observer/stats/wrapper/conn_test.go
+++ b/observer/stats/wrapper/conn_test.go
@@ -306,6 +306,41 @@ func TestWrapUDPConn_Nil(t *testing.T) {
 	}
 }
 
+func TestWrapUDPConn_Conns(t *testing.T) {
+	st := ostats.NewStats(false)
+	WrapUDPConn(&fakePacketConn{}, st)
+
+	if v := st.Get(stats.KindTotalConns); v != 1 {
+		t.Errorf("totalConns = %d, want 1", v)
+	}
+	if v := st.Get(stats.KindCurrentConns); v != 1 {
+		t.Errorf("currentConns = %d, want 1", v)
+	}
+}
+
+func TestWrapUDPConn_Close(t *testing.T) {
+	st := ostats.NewStats(false)
+	wrapped := WrapUDPConn(&fakePacketConn{}, st)
+
+	if err := wrapped.Close(); err != nil {
+		t.Fatalf("close: %v", err)
+	}
+	if v := st.Get(stats.KindCurrentConns); v != 0 {
+		t.Errorf("currentConns after close = %d, want 0", v)
+	}
+
+	// A second Close must be a no-op: currentConns must not go negative.
+	if err := wrapped.Close(); err != nil {
+		t.Fatalf("second close: %v", err)
+	}
+	if v := st.Get(stats.KindCurrentConns); v != 0 {
+		t.Errorf("currentConns after second close = %d, want 0", v)
+	}
+	if v := st.Get(stats.KindTotalConns); v != 1 {
+		t.Errorf("totalConns = %d, want 1 (monotonic)", v)
+	}
+}
+
 func TestWrapUDPConn_ReadFrom(t *testing.T) {
 	st := ostats.NewStats(false)
 	fpc := &fakePacketConn{readBuf: []byte("udp-packet")}