counting.go 596 B

123456789101112131415161718192021222324252627282930313233343536
  1. package protocol
  2. import (
  3. "io"
  4. "sync/atomic"
  5. )
  6. type countingReader struct {
  7. io.Reader
  8. tot uint64
  9. }
  10. func (c *countingReader) Read(bs []byte) (int, error) {
  11. n, err := c.Reader.Read(bs)
  12. atomic.AddUint64(&c.tot, uint64(n))
  13. return n, err
  14. }
  15. func (c *countingReader) Tot() uint64 {
  16. return atomic.LoadUint64(&c.tot)
  17. }
  18. type countingWriter struct {
  19. io.Writer
  20. tot uint64
  21. }
  22. func (c *countingWriter) Write(bs []byte) (int, error) {
  23. n, err := c.Writer.Write(bs)
  24. atomic.AddUint64(&c.tot, uint64(n))
  25. return n, err
  26. }
  27. func (c *countingWriter) Tot() uint64 {
  28. return atomic.LoadUint64(&c.tot)
  29. }