simple metrics to capture packet data size both compressed & uncompressed

Author: dvilaverde

packet/conn.go

@@ -67,6 +67,8 @@ type Conn struct {
 	compressedHeader [7]byte
 
 	compressedReader io.Reader
+
+	Stats *Stats
 }
 
 func NewConn(conn net.Conn) *Conn {
@@ -274,6 +276,9 @@ func (c *Conn) ReadPacketTo(w io.Writer) error {
 // will modify data inplace
 func (c *Conn) WritePacket(data []byte) error {
 	length := len(data) - 4
+	if c.Stats != nil {
+		c.Stats.AddTxUncompressedSize(uint64(length))
+	}
 
 	for length >= MaxPayloadLen {
 		data[0] = 0xff
@@ -378,6 +383,9 @@ func (c *Conn) writeCompressed(data []byte) (n int, err error) {
 		return 0, err
 	}
 
+	if c.Stats != nil {
+		c.Stats.AddTxCompressedSize(uint64(compressedPacket.Len()))
+	}
 	_, err = c.Write(compressedPacket.Bytes())
 	if err != nil {
 		return 0, err

packet/metrics.go

@@ -0,0 +1,24 @@
+package packet
+
+import "sync/atomic"
+
+type Stats struct {
+	packetTxCompressedSize   atomic.Uint64
+	packetTxUncompressedSize atomic.Uint64
+}
+
+func (p *Stats) AddTxCompressedSize(size uint64) {
+	p.packetTxCompressedSize.Add(size)
+}
+
+func (p *Stats) AddTxUncompressedSize(size uint64) {
+	p.packetTxUncompressedSize.Add(size)
+}
+
+func (p *Stats) GetTxCompressedSize() uint64 {
+	return p.packetTxCompressedSize.Load()
+}
+
+func (p *Stats) GetTxUncompressedSize() uint64 {
+	return p.packetTxUncompressedSize.Load()
+}