diff --git a/countrw/cr.go b/countrw/cr.go new file mode 100644 index 00000000..817dd255 --- /dev/null +++ b/countrw/cr.go @@ -0,0 +1,26 @@ +package countrw + +import ( + "io" + "sync/atomic" +) + +type CountReader struct { + r io.Reader + c *uint64 +} + +func NewCountReader(r io.Reader) *CountReader { + return &CountReader{r: r, c: new(uint64)} +} + +func (c *CountReader) Reset() { atomic.StoreUint64(c.c, 0) } +func (c *CountReader) Count() uint64 { return atomic.LoadUint64(c.c) } + +func (c *CountReader) Read(p []byte) (int, error) { + n, err := c.r.Read(p) + atomic.AddUint64(c.c, uint64(n)) + return n, err +} + +var _ io.Reader = (*CountReader)(nil) diff --git a/countrw/cw.go b/countrw/cw.go new file mode 100644 index 00000000..5d27d9c2 --- /dev/null +++ b/countrw/cw.go @@ -0,0 +1,26 @@ +package countrw + +import ( + "io" + "sync/atomic" +) + +type CountWriter struct { + r io.Writer + c *uint64 +} + +func NewCountWriter(r io.Writer) *CountWriter { + return &CountWriter{r: r, c: new(uint64)} +} + +func (c *CountWriter) Reset() { atomic.StoreUint64(c.c, 0) } +func (c *CountWriter) Count() uint64 { return atomic.LoadUint64(c.c) } + +func (c *CountWriter) Write(p []byte) (int, error) { + n, err := c.r.Write(p) + atomic.AddUint64(c.c, uint64(n)) + return n, err +} + +var _ io.Writer = (*CountWriter)(nil)