github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/netx/stat.go (about)

     1  package netx
     2  
     3  import (
     4  	"net"
     5  	"sync/atomic"
     6  )
     7  
     8  // StatConn statistics read and write bytes for net connection.
     9  type StatConn struct {
    10  	net.Conn
    11  	Reads, Writes *uint64
    12  }
    13  
    14  func NewStatConn(c net.Conn) *StatConn {
    15  	var r, w uint64
    16  	return &StatConn{Conn: c, Reads: &r, Writes: &w}
    17  }
    18  
    19  func NewStatConnReadWrite(c net.Conn, reads, writes *uint64) *StatConn {
    20  	return &StatConn{Conn: c, Reads: reads, Writes: writes}
    21  }
    22  
    23  // Read reads bytes from net connection.
    24  func (s *StatConn) Read(b []byte) (n int, err error) {
    25  	if n, err = s.Conn.Read(b); err == nil {
    26  		atomic.AddUint64(s.Reads, uint64(n))
    27  	}
    28  
    29  	return
    30  }
    31  
    32  // Write writes bytes to net.
    33  func (s *StatConn) Write(b []byte) (n int, err error) {
    34  	if n, err = s.Conn.Write(b); err == nil {
    35  		atomic.AddUint64(s.Writes, uint64(n))
    36  	}
    37  
    38  	return
    39  }
    40  
    41  func (s StatConn) ReadBytes() uint64  { return *s.Reads }
    42  func (s StatConn) WriteBytes() uint64 { return *s.Writes }