counting.go 1.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. // Copyright (C) 2014 The Protocol Authors.
  2. package protocol
  3. import (
  4. "io"
  5. "sync/atomic"
  6. "time"
  7. )
  8. type countingReader struct {
  9. io.Reader
  10. tot atomic.Int64 // bytes
  11. last atomic.Int64 // unix nanos
  12. }
  13. var (
  14. totalIncoming atomic.Int64
  15. totalOutgoing atomic.Int64
  16. )
  17. func (c *countingReader) Read(bs []byte) (int, error) {
  18. n, err := c.Reader.Read(bs)
  19. c.tot.Add(int64(n))
  20. totalIncoming.Add(int64(n))
  21. c.last.Store(time.Now().UnixNano())
  22. return n, err
  23. }
  24. func (c *countingReader) Tot() int64 { return c.tot.Load() }
  25. func (c *countingReader) Last() time.Time {
  26. return time.Unix(0, c.last.Load())
  27. }
  28. type countingWriter struct {
  29. io.Writer
  30. tot atomic.Int64 // bytes
  31. last atomic.Int64 // unix nanos
  32. }
  33. func (c *countingWriter) Write(bs []byte) (int, error) {
  34. n, err := c.Writer.Write(bs)
  35. c.tot.Add(int64(n))
  36. totalOutgoing.Add(int64(n))
  37. c.last.Store(time.Now().UnixNano())
  38. return n, err
  39. }
  40. func (c *countingWriter) Tot() int64 { return c.tot.Load() }
  41. func (c *countingWriter) Last() time.Time {
  42. return time.Unix(0, c.last.Load())
  43. }
  44. func TotalInOut() (int64, int64) {
  45. return totalIncoming.Load(), totalOutgoing.Load()
  46. }