counting.go 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. // Copyright (C) 2014 Jakob Borg and Contributors (see the CONTRIBUTORS file).
  2. // All rights reserved. Use of this source code is governed by an MIT-style
  3. // license that can be found in the LICENSE file.
  4. package protocol
  5. import (
  6. "io"
  7. "sync/atomic"
  8. "time"
  9. )
  10. type countingReader struct {
  11. io.Reader
  12. tot uint64 // bytes
  13. last int64 // unix nanos
  14. }
  15. var (
  16. totalIncoming uint64
  17. totalOutgoing uint64
  18. )
  19. func (c *countingReader) Read(bs []byte) (int, error) {
  20. n, err := c.Reader.Read(bs)
  21. atomic.AddUint64(&c.tot, uint64(n))
  22. atomic.AddUint64(&totalIncoming, uint64(n))
  23. atomic.StoreInt64(&c.last, time.Now().UnixNano())
  24. return n, err
  25. }
  26. func (c *countingReader) Tot() uint64 {
  27. return atomic.LoadUint64(&c.tot)
  28. }
  29. func (c *countingReader) Last() time.Time {
  30. return time.Unix(0, atomic.LoadInt64(&c.last))
  31. }
  32. type countingWriter struct {
  33. io.Writer
  34. tot uint64 // bytes
  35. last int64 // unix nanos
  36. }
  37. func (c *countingWriter) Write(bs []byte) (int, error) {
  38. n, err := c.Writer.Write(bs)
  39. atomic.AddUint64(&c.tot, uint64(n))
  40. atomic.AddUint64(&totalOutgoing, uint64(n))
  41. atomic.StoreInt64(&c.last, time.Now().UnixNano())
  42. return n, err
  43. }
  44. func (c *countingWriter) Tot() uint64 {
  45. return atomic.LoadUint64(&c.tot)
  46. }
  47. func (c *countingWriter) Last() time.Time {
  48. return time.Unix(0, atomic.LoadInt64(&c.last))
  49. }
  50. func TotalInOut() (uint64, uint64) {
  51. return atomic.LoadUint64(&totalIncoming), atomic.LoadUint64(&totalOutgoing)
  52. }