feat(countrw): add countrw package

This commit is contained in:
sentriz
2022-03-10 00:42:50 +00:00
committed by Senan Kelly
parent c7862761bd
commit 5155dee2e8
2 changed files with 52 additions and 0 deletions

26
countrw/cr.go Normal file
View File

@@ -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)

26
countrw/cw.go Normal file
View File

@@ -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)