counting.go 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  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 int64 // bytes
  11. last int64 // unix nanos
  12. }
  13. var (
  14. totalIncoming int64
  15. totalOutgoing int64
  16. )
  17. func (c *countingReader) Read(bs []byte) (int, error) {
  18. n, err := c.Reader.Read(bs)
  19. atomic.AddInt64(&c.tot, int64(n))
  20. atomic.AddInt64(&totalIncoming, int64(n))
  21. atomic.StoreInt64(&c.last, time.Now().UnixNano())
  22. return n, err
  23. }
  24. func (c *countingReader) Tot() int64 {
  25. return atomic.LoadInt64(&c.tot)
  26. }
  27. func (c *countingReader) Last() time.Time {
  28. return time.Unix(0, atomic.LoadInt64(&c.last))
  29. }
  30. type countingWriter struct {
  31. io.Writer
  32. tot int64 // bytes
  33. last int64 // unix nanos
  34. }
  35. func (c *countingWriter) Write(bs []byte) (int, error) {
  36. n, err := c.Writer.Write(bs)
  37. atomic.AddInt64(&c.tot, int64(n))
  38. atomic.AddInt64(&totalOutgoing, int64(n))
  39. atomic.StoreInt64(&c.last, time.Now().UnixNano())
  40. return n, err
  41. }
  42. func (c *countingWriter) Tot() int64 {
  43. return atomic.LoadInt64(&c.tot)
  44. }
  45. func (c *countingWriter) Last() time.Time {
  46. return time.Unix(0, atomic.LoadInt64(&c.last))
  47. }
  48. func TotalInOut() (int64, int64) {
  49. return atomic.LoadInt64(&totalIncoming), atomic.LoadInt64(&totalOutgoing)
  50. }