2020-04-10 21:58:28 +02:00
|
|
|
package bytecounter
|
|
|
|
|
|
|
|
import (
|
|
|
|
"io"
|
|
|
|
"sync/atomic"
|
|
|
|
)
|
|
|
|
|
|
|
|
// ReadCloser wraps an io.ReadCloser, reimplementing
|
|
|
|
// its interface and counting the bytes written to during copying.
|
|
|
|
type ReadCloser interface {
|
|
|
|
io.ReadCloser
|
2021-08-16 10:11:37 +02:00
|
|
|
Count() uint64
|
2020-04-10 21:58:28 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
// NewReadCloser wraps rc.
|
|
|
|
func NewReadCloser(rc io.ReadCloser) ReadCloser {
|
|
|
|
return &readCloser{rc, 0}
|
|
|
|
}
|
|
|
|
|
|
|
|
type readCloser struct {
|
|
|
|
rc io.ReadCloser
|
2021-08-16 10:11:37 +02:00
|
|
|
count uint64
|
2020-04-10 21:58:28 +02:00
|
|
|
}
|
|
|
|
|
2021-08-16 10:11:37 +02:00
|
|
|
func (r *readCloser) Count() uint64 {
|
|
|
|
return atomic.LoadUint64(&r.count)
|
2020-04-10 21:58:28 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
var _ io.ReadCloser = &readCloser{}
|
|
|
|
|
|
|
|
func (r *readCloser) Close() error {
|
|
|
|
return r.rc.Close()
|
|
|
|
}
|
|
|
|
|
|
|
|
func (r *readCloser) Read(p []byte) (int, error) {
|
|
|
|
n, err := r.rc.Read(p)
|
2021-08-16 10:11:37 +02:00
|
|
|
if n < 0 {
|
|
|
|
panic("expecting n >= 0")
|
|
|
|
}
|
|
|
|
atomic.AddUint64(&r.count, uint64(n))
|
2020-04-10 21:58:28 +02:00
|
|
|
return n, err
|
|
|
|
}
|