1
0

counting.go 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. // Copyright (C) 2014 Jakob Borg and Contributors (see the CONTRIBUTORS file).
  2. //
  3. // This program is free software: you can redistribute it and/or modify it
  4. // under the terms of the GNU General Public License as published by the Free
  5. // Software Foundation, either version 3 of the License, or (at your option)
  6. // any later version.
  7. //
  8. // This program is distributed in the hope that it will be useful, but WITHOUT
  9. // ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  10. // FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
  11. // more details.
  12. //
  13. // You should have received a copy of the GNU General Public License along
  14. // with this program. If not, see <http://www.gnu.org/licenses/>.
  15. package protocol
  16. import (
  17. "io"
  18. "sync/atomic"
  19. "time"
  20. )
  21. type countingReader struct {
  22. io.Reader
  23. tot uint64 // bytes
  24. last int64 // unix nanos
  25. }
  26. var (
  27. totalIncoming uint64
  28. totalOutgoing uint64
  29. )
  30. func (c *countingReader) Read(bs []byte) (int, error) {
  31. n, err := c.Reader.Read(bs)
  32. atomic.AddUint64(&c.tot, uint64(n))
  33. atomic.AddUint64(&totalIncoming, uint64(n))
  34. atomic.StoreInt64(&c.last, time.Now().UnixNano())
  35. return n, err
  36. }
  37. func (c *countingReader) Tot() uint64 {
  38. return atomic.LoadUint64(&c.tot)
  39. }
  40. func (c *countingReader) Last() time.Time {
  41. return time.Unix(0, atomic.LoadInt64(&c.last))
  42. }
  43. type countingWriter struct {
  44. io.Writer
  45. tot uint64 // bytes
  46. last int64 // unix nanos
  47. }
  48. func (c *countingWriter) Write(bs []byte) (int, error) {
  49. n, err := c.Writer.Write(bs)
  50. atomic.AddUint64(&c.tot, uint64(n))
  51. atomic.AddUint64(&totalOutgoing, uint64(n))
  52. atomic.StoreInt64(&c.last, time.Now().UnixNano())
  53. return n, err
  54. }
  55. func (c *countingWriter) Tot() uint64 {
  56. return atomic.LoadUint64(&c.tot)
  57. }
  58. func (c *countingWriter) Last() time.Time {
  59. return time.Unix(0, atomic.LoadInt64(&c.last))
  60. }
  61. func TotalInOut() (uint64, uint64) {
  62. return atomic.LoadUint64(&totalIncoming), atomic.LoadUint64(&totalOutgoing)
  63. }