counting.go 856 B

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647
  1. package protocol
  2. import (
  3. "io"
  4. "sync/atomic"
  5. )
  6. type countingReader struct {
  7. io.Reader
  8. tot uint64
  9. }
  10. var (
  11. totalIncoming uint64
  12. totalOutgoing uint64
  13. )
  14. func (c *countingReader) Read(bs []byte) (int, error) {
  15. n, err := c.Reader.Read(bs)
  16. atomic.AddUint64(&c.tot, uint64(n))
  17. atomic.AddUint64(&totalIncoming, uint64(n))
  18. return n, err
  19. }
  20. func (c *countingReader) Tot() uint64 {
  21. return atomic.LoadUint64(&c.tot)
  22. }
  23. type countingWriter struct {
  24. io.Writer
  25. tot uint64
  26. }
  27. func (c *countingWriter) Write(bs []byte) (int, error) {
  28. n, err := c.Writer.Write(bs)
  29. atomic.AddUint64(&c.tot, uint64(n))
  30. atomic.AddUint64(&totalOutgoing, uint64(n))
  31. return n, err
  32. }
  33. func (c *countingWriter) Tot() uint64 {
  34. return atomic.LoadUint64(&c.tot)
  35. }
  36. func TotalInOut() (uint64, uint64) {
  37. return atomic.LoadUint64(&totalIncoming), atomic.LoadUint64(&totalOutgoing)
  38. }