counting.go 1.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. // Copyright (C) 2014 Jakob Borg and other contributors. All rights reserved.
  2. // Use of this source code is governed by an MIT-style license that can be
  3. // found in the LICENSE file.
  4. package protocol
  5. import (
  6. "io"
  7. "sync/atomic"
  8. )
  9. type countingReader struct {
  10. io.Reader
  11. tot uint64
  12. }
  13. var (
  14. totalIncoming uint64
  15. totalOutgoing uint64
  16. )
  17. func (c *countingReader) Read(bs []byte) (int, error) {
  18. n, err := c.Reader.Read(bs)
  19. atomic.AddUint64(&c.tot, uint64(n))
  20. atomic.AddUint64(&totalIncoming, uint64(n))
  21. return n, err
  22. }
  23. func (c *countingReader) Tot() uint64 {
  24. return atomic.LoadUint64(&c.tot)
  25. }
  26. type countingWriter struct {
  27. io.Writer
  28. tot uint64
  29. }
  30. func (c *countingWriter) Write(bs []byte) (int, error) {
  31. n, err := c.Writer.Write(bs)
  32. atomic.AddUint64(&c.tot, uint64(n))
  33. atomic.AddUint64(&totalOutgoing, uint64(n))
  34. return n, err
  35. }
  36. func (c *countingWriter) Tot() uint64 {
  37. return atomic.LoadUint64(&c.tot)
  38. }
  39. func TotalInOut() (uint64, uint64) {
  40. return atomic.LoadUint64(&totalIncoming), atomic.LoadUint64(&totalOutgoing)
  41. }