feat(observer/stats): track UDP connection counts in WrapUDPConn
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.
This commit is contained in:
@@ -146,23 +146,50 @@ func (c *packetConn) Context() context.Context {
|
|||||||
|
|
||||||
type udpConn struct {
|
type udpConn struct {
|
||||||
net.PacketConn
|
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
|
// WrapUDPConn wraps a net.PacketConn as a udp.Conn to track connection and
|
||||||
// bytes across all read/write methods (Read, Write, ReadFrom, WriteTo,
|
// traffic statistics. It increments total and current connection counts (one
|
||||||
// ReadFromUDP, ReadMsgUDP, WriteToUDP, WriteMsgUDP). If pc or stats is nil,
|
// 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.
|
// it returns nil.
|
||||||
func WrapUDPConn(pc net.PacketConn, stats stats.Stats) udp.Conn {
|
func WrapUDPConn(pc net.PacketConn, pStats stats.Stats) udp.Conn {
|
||||||
if pc == nil || stats == nil {
|
if pc == nil || pStats == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pStats.Add(stats.KindTotalConns, 1)
|
||||||
|
pStats.Add(stats.KindCurrentConns, 1)
|
||||||
|
|
||||||
return &udpConn{
|
return &udpConn{
|
||||||
PacketConn: pc,
|
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 {
|
func (c *udpConn) RemoteAddr() net.Addr {
|
||||||
if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok {
|
if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok {
|
||||||
return nc.RemoteAddr()
|
return nc.RemoteAddr()
|
||||||
|
|||||||
@@ -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) {
|
func TestWrapUDPConn_ReadFrom(t *testing.T) {
|
||||||
st := ostats.NewStats(false)
|
st := ostats.NewStats(false)
|
||||||
fpc := &fakePacketConn{readBuf: []byte("udp-packet")}
|
fpc := &fakePacketConn{readBuf: []byte("udp-packet")}
|
||||||
|
|||||||
Reference in New Issue
Block a user